Home
last modified time | relevance | path

Searched refs:InputXlaShape (Results 1 – 8 of 8) sorted by relevance

/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dspmd_manual_sharding_ops.cc40 auto input_shape_or = ctx->InputXlaShape(0); in Compile()
100 auto input_shape_or = ctx->InputXlaShape(0); in Compile()
Dreshape_op.cc113 auto input_xla_shape = ctx->InputXlaShape(0); in Compile()
174 VLOG(2) << "Reshape from " << ctx->InputXlaShape(0)->ToString() in Compile()
Dshape_op.cc104 ctx->InputXlaShape("input").ValueOrDie(), "", false, {}, in Compile()
382 auto input_shape = ctx->InputXlaShape(0).ValueOrDie(); in Compile()
Dsegment_reduction_ops.cc105 ctx->InputXlaShape(0)->is_dynamic_dimension(i)); in Compile()
Dstrided_slice_op.cc140 xla::Shape input_xla_shape = ctx->InputXlaShape(0).ValueOrDie(); in EmitDynamicSlice()
414 xla::Shape grad_shape = ctx->InputXlaShape(4).ValueOrDie(); in CompileAsDynamicUpdateSlice()
/external/tensorflow/tensorflow/compiler/tf2xla/
Dxla_op_kernel.h106 xla::StatusOr<xla::Shape> InputXlaShape(int index);
107 xla::StatusOr<xla::Shape> InputXlaShape(absl::string_view name);
Dxla_op_kernel.cc77 xla::StatusOr<xla::Shape> XlaOpKernelContext::InputXlaShape(int index) { in InputXlaShape() function in tensorflow::XlaOpKernelContext
81 xla::StatusOr<xla::Shape> XlaOpKernelContext::InputXlaShape( in InputXlaShape() function in tensorflow::XlaOpKernelContext
/external/tensorflow/tensorflow/core/tpu/kernels/xla/
Dsegment_reduction_ops.cc95 ctx->InputXlaShape(0)->is_dynamic_dimension(i)); in Compile()