Searched refs:InputXlaShape (Results 1 – 15 of 15) sorted by relevance
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | cwise_ops.cc | 39 xla::Shape lhs_xla_shape = ctx->InputXlaShape(0).ValueOrDie(); in Compile()
40 xla::Shape rhs_xla_shape = ctx->InputXlaShape(1).ValueOrDie(); in Compile()
|
| D | reshape_op.cc | 40 auto input_xla_shape = ctx->InputXlaShape(0); in Compile()
192 VLOG(2) << "Reshape from " << ctx->InputXlaShape(0)->ToString() in Compile()
|
| D | dynamic_partition_op.cc | 111 xla::Shape data_shape = ctx->InputXlaShape(0).value(); in Compile()
112 xla::Shape partition_shape = ctx->InputXlaShape(1).value(); in Compile()
|
| D | to_bool_op.cc | 41 TF_ASSIGN_OR_RETURN(auto shape, ctx->InputXlaShape(0)); in DoCompile()
|
| D | xla_call_module_op.cc | 161 OP_REQUIRES_VALUE(xla::Shape xla_shape, ctx, ctx->InputXlaShape(i)); in RefineDynamicShapes()
219 OP_REQUIRES_VALUE(xla::Shape xla_shape, ctx, ctx->InputXlaShape(arg_idx)); in PopulateDimArgInputs()
|
| D | spmd_manual_sharding_ops.cc | 91 auto status_or_input_shape = ctx->InputXlaShape(0); in Compile()
|
| D | shape_op.cc | 104 ctx->InputXlaShape("input").ValueOrDie(), "", false, {}, in Compile()
416 auto input_shape = ctx->InputXlaShape(0).ValueOrDie(); in Compile()
|
| D | approx_topk_op.cc | 67 xla::Shape op_shape = ctx->InputXlaShape(0).value(); in Compile()
|
| D | segment_reduction_ops.cc | 107 ctx->InputXlaShape(0)->is_dynamic_dimension(i)); in Compile()
|
| D | strided_slice_op.cc | 147 xla::Shape input_xla_shape = ctx->InputXlaShape(0).ValueOrDie(); in EmitDynamicSlice()
451 xla::Shape grad_shape = ctx->InputXlaShape(4).ValueOrDie(); in CompileAsDynamicUpdateSlice()
|
| D | light_outside_compilation.cc | 106 TF_ASSIGN_OR_RETURN(xla::Shape xla_shape, ctx->InputXlaShape(i)); in CompileToCustomCallCallingTfKernel()
|
| /external/tensorflow/tensorflow/compiler/tf2xla/ |
| D | xla_op_kernel.h | 115 StatusOr<xla::Shape> InputXlaShape(int index);
116 StatusOr<xla::Shape> InputXlaShape(absl::string_view name);
|
| D | xla_op_kernel.cc | 87 StatusOr<xla::Shape> XlaOpKernelContext::InputXlaShape(int index) { in InputXlaShape() function in tensorflow::XlaOpKernelContext
91 StatusOr<xla::Shape> XlaOpKernelContext::InputXlaShape(absl::string_view name) { in InputXlaShape() function in tensorflow::XlaOpKernelContext
142 if (this->InputXlaShape(index)->is_dynamic()) { in ConstantInput()
146 this->InputXlaShape(index)->ToString()); in ConstantInput()
|
| D | mlir_xla_op_kernel.cc | 46 arg.shape = ctx->InputXlaShape(i).ValueOrDie(); in ContextToXlaArgs()
|
| /external/tensorflow/tensorflow/core/tpu/kernels/xla/ |
| D | segment_reduction_ops.cc | 97 ctx->InputXlaShape(0)->is_dynamic_dimension(i)); in Compile()
|