コード例 #1
0
 def applicable_jax2tf_limitation(l: Jax2TfLimitation) -> bool:
   # The CheckShapePolymorphism uses tf.function, so we care about "graph"
   return l.filter(device=device, dtype=h.dtype, mode="graph")
コード例 #2
0
ファイル: shape_poly_test.py プロジェクト: memari-majid/jax 
 def applicable_jax2tf_limitation(l: Jax2TfLimitation) -> bool:
     return (l.filter(device=device, dtype=h.dtype, mode="graph")
             and l.expect_tf_error)