コード例 #1
0
ファイル: djax_test.py プロジェクト: yashk2810/jax 
 def f(x):
   return nonzero(x)
コード例 #2
0
ファイル: djax_test.py プロジェクト: yashk2810/jax 
 def f(x):
   return reduce_sum(nonzero(x), tuple(range(len(x.shape))))
コード例 #3
0
ファイル: djax_test.py プロジェクト: yashk2810/jax 
 def f(x):
   nonzero_idx = nonzero(x)
   return reduce_sum(nonzero_idx)
コード例 #4
0
ファイル: djax_test.py プロジェクト: yashk2810/jax 
 def f(x, n):
   y = nonzero(x)
   return broadcast(y, n)