非弾性構成人工ニューラルネットワークのための一般化された双対ポテンシャル：有限ひずみにおけるJAX実装