Home
last modified time | relevance | path

Searched refs:xla_shape (Results 1 – 4 of 4) sorted by relevance

/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Drandom_ops.cc41 xla::Shape xla_shape; in Compile() local
42 OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(dtype, shape, &xla_shape)); in Compile()
46 XlaHelpers::Zero(b, dtype), XlaHelpers::One(b, dtype), xla_shape); in Compile()
65 xla::Shape xla_shape; in Compile() local
67 TensorShapeToXLAShape(input_type(1), shape, &xla_shape)); in Compile()
80 ctx->SetOutput(0, ctx->builder()->RngUniform(minval, maxval, xla_shape)); in Compile()
100 xla::Shape xla_shape; in Compile() local
101 OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(dtype, shape, &xla_shape)); in Compile()
107 XlaHelpers::Zero(b, dtype), XlaHelpers::One(b, dtype), xla_shape); in Compile()
128 xla::Shape xla_shape; in Compile() local
[all …]
Dindex_ops_cpu.cc89 xla::Shape xla_shape = in Compile() local
97 output = b.CustomCall("argmax_float_1d_xla_impl", args, xla_shape); in Compile()
100 output = b.CustomCall("argmax_float_2d_xla_impl", args, xla_shape); in Compile()
/external/tensorflow/tensorflow/compiler/tf2xla/
Dxla_compiler.cc233 xla::Shape* xla_shape) { in XLAShapeForArgument() argument
237 xla_shape); in XLAShapeForArgument()
239 return TensorShapeToXLAShape(arg.type, arg.shape, xla_shape); in XLAShapeForArgument()
248 xla_shape); in XLAShapeForArgument()
258 TF_RETURN_IF_ERROR(TensorShapeToXLAShape(arg.type, shape, xla_shape)); in XLAShapeForArgument()
262 arg.tensor_array_gradients.size() + 1, *xla_shape); in XLAShapeForArgument()
263 *xla_shape = xla::ShapeUtil::MakeTupleShape(tuple_shape); in XLAShapeForArgument()
278 *xla_shape = xla::ShapeUtil::MakeTupleShape( in XLAShapeForArgument()
Dxla_compiler.h290 Status XLAShapeForArgument(const Argument& arg, xla::Shape* xla_shape);