Searched refs:dim_t (Results 1 – 4 of 4) sorted by relevance
627 Tensor dim_t; in TEST() local628 op.input_tensors[1] = &dim_t; in TEST()632 dim_t = test::AsScalar<int32>(idx); in TEST()639 dim_t = test::AsScalar<int32>(idx); in TEST()644 dim_t = test::AsScalar<int64>(idx); in TEST()649 dim_t = test::AsScalar<int32>(idx); in TEST()654 dim_t = test::AsScalar<int64>(idx); in TEST()661 dim_t = test::AsScalar<int32>(idx); in TEST()666 dim_t = test::AsScalar<int64>(idx); in TEST()672 dim_t = test::AsScalar<int32>(idx); in TEST()[all …]
1060 const Tensor* dim_t = c->input_tensor(1); in ArgOpShape() local1061 if (dim_t == nullptr) { in ArgOpShape()1075 if (dim_t->dtype() == DT_INT32) { in ArgOpShape()1076 dimension_val = dim_t->scalar<int32>()(); in ArgOpShape()1078 dimension_val = dim_t->scalar<int64>()(); in ArgOpShape()2074 const Tensor* dim_t = c->input_tensor(0); in __anon57f5a5dd1c02() local2077 int32 dim = dim_t == nullptr ? InferenceContext::kUnknownDim in __anon57f5a5dd1c02()2078 : dim_t->scalar<int32>()(); in __anon57f5a5dd1c02()
2084 const Tensor* dim_t = c->input_tensor(1); in __anondb9326b22f02() local2085 if (dim_t != nullptr && dim_t->NumElements() != 1) { in __anondb9326b22f02()2089 if (dim_t == nullptr || !c->RankKnown(input)) { in __anondb9326b22f02()2095 if (dim_t->dtype() == DT_INT32) { in __anondb9326b22f02()2096 dim = static_cast<int64>(dim_t->flat<int32>()(0)); in __anondb9326b22f02()2098 dim = dim_t->flat<int64>()(0); in __anondb9326b22f02()
147 const Tensor& dim_t = ctx->input(1); in Compute() local149 ctx, (dim_t.NumElements() == 1), in Compute()151 DCHECK_EQ(dim_t.dtype(), DataTypeToEnum<Tdim>::v()); in Compute()152 Tdim dim = *static_cast<const Tdim*>(DMAHelper::base(&dim_t)); in Compute()