H A D | core_named_ops.py | 6 Batch = S.Batch variable 468 lhs=TensorDef(TV.LhsType, Batch, S.M, S.K, S.M0, S.K0), 469 rhs=TensorDef(TV.RhsType, Batch, S.N, S.K, S.N0, S.K0), 470 accum=TensorDef(TV.AccumType, Batch, S.M, S.N, S.M0, S.N0, output=True), 475 Besides the outermost batch dimension has the same semantic as 477 non-batch dimensions are the same as linalg.mmt4d vs. linalg.matmul. See the 489 A=TensorDef(T1, Batch, S.M, S.K), 490 B=TensorDef(T2, Batch, S.K, S.N), 491 C=TensorDef(U, Batch, argument 512 batch_matmul( A=TensorDef(T1, Batch, S.M, S.K), B=TensorDef(T2, Batch, S.K, S.N), C=TensorDef(U, Batch, S.M, S.N, output=True), ) global() argument 530 batch_matmul_transpose_a( A=TensorDef(T1, Batch, S.K, S.M), B=TensorDef(T2, Batch, S.K, S.N), C=TensorDef(U, Batch, S.M, S.N, output=True), ) global() argument 549 batch_matmul_transpose_b( A=TensorDef(T1, Batch, S.M, S.K), B=TensorDef(T2, Batch, S.N, S.K), C=TensorDef(U, Batch, S.M, S.N, output=True), ) global() argument 568 quantized_batch_matmul( A=TensorDef(T1, Batch, S.M, S.K), B=TensorDef(T2, Batch, S.K, S.N), AZp=ScalarDef(I32), BZp=ScalarDef(I32), C=TensorDef(U, Batch, S.M, S.N, output=True), ) global() argument 589 batch_reduce_matmul( A=TensorDef(T1, Batch, S.M, S.K), B=TensorDef(T2, Batch, S.K, S.N), C=TensorDef(U, S.M, S.N, output=True), ) global() argument 636 batch_matvec( A=TensorDef(T1, Batch, S.M, S.K), B=TensorDef(T2, Batch, S.K), C=TensorDef(U, Batch, S.M, output=True), ) global() argument 654 batch_vecmat( A=TensorDef(T1, Batch, S.K), B=TensorDef(T2, Batch, S.K, S.N), C=TensorDef(U, Batch, S.N, output=True), ) global() argument [all...] |