Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dstateless_random_ops.cc87 xla::Shape xla_shape; in Compile() local
88 OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(DT_FLOAT, shape, &xla_shape)); in Compile()
90 xla_shape, dtype_, seed, xla::ConstantR0<float>(builder, 0.0), in Compile()
137 xla::Shape xla_shape; in Compile() local
138 OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(dtype_, shape, &xla_shape)); in Compile()
140 StatelessRandomUniformImpl(xla_shape, dtype_, seed, minval, maxval); in Compile()
174 xla::Shape xla_shape; in Compile() local
175 OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(DT_FLOAT, shape, &xla_shape)); in Compile()
177 xla_shape, dtype_, seed, in Compile()
216 xla::Shape xla_shape; in Compile() local
[all …]
Drandom_ops.cc47 xla::Shape xla_shape; in Compile() local
48 OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(dtype, shape, &xla_shape)); in Compile()
52 XlaHelpers::One(b, dtype), xla_shape); in Compile()
215 xla::Shape xla_shape; in Compile() local
217 TensorShapeToXLAShape(input_type(1), shape, &xla_shape)); in Compile()
230 ctx->SetOutput(0, xla::RngUniform(minval, maxval, xla_shape)); in Compile()
250 xla::Shape xla_shape; in Compile() local
251 OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(dtype, shape, &xla_shape)); in Compile()
257 XlaHelpers::One(b, dtype), xla_shape); in Compile()
278 xla::Shape xla_shape; in Compile() local
[all …]
Dstateful_random_ops.cc259 xla::Shape xla_shape; in Compile() local
260 TF_RETURN_IF_ERROR(TensorShapeToXLAShape(DT_FLOAT, shape, &xla_shape)); in Compile()
263 key, counter, xla_shape, in Compile()
306 xla::Shape xla_shape; in Compile() local
307 TF_RETURN_IF_ERROR(TensorShapeToXLAShape(dtype_, shape, &xla_shape)); in Compile()
308 return StatefulRngUniform(key, counter, xla_shape, minval, maxval); in Compile()
339 xla::Shape xla_shape; in Compile() local
340 TF_RETURN_IF_ERROR(TensorShapeToXLAShape(dtype_, shape, &xla_shape)); in Compile()
341 return StatefulRngUniformFullInt(key, counter, xla_shape); in Compile()
Dindex_ops_cpu.cc106 xla::Shape xla_shape = xla::ShapeUtil::MakeShapeWithDescendingLayout( in Compile() local
122 xla_shape, arg_shapes); in Compile()
125 xla_shape, arg_shapes); in Compile()
/external/tensorflow/tensorflow/compiler/tf2xla/
Dliteral_util.cc27 xla::Shape xla_shape; in HostTensorToBorrowingLiteral() local
29 host_tensor.shape(), &xla_shape)); in HostTensorToBorrowingLiteral()
31 static_cast<const char*>(DMAHelper::base(&host_tensor)), xla_shape); in HostTensorToBorrowingLiteral()
43 xla::Shape xla_shape; in HostTensorToMutableBorrowingLiteral() local
45 host_tensor->shape(), &xla_shape)); in HostTensorToMutableBorrowingLiteral()
46 return HostTensorToMutableBorrowingLiteral(xla_shape, host_tensor, literal); in HostTensorToMutableBorrowingLiteral()
50 const xla::Shape& xla_shape, Tensor* host_tensor, in HostTensorToMutableBorrowingLiteral() argument
53 static_cast<const char*>(DMAHelper::base(host_tensor)), xla_shape); in HostTensorToMutableBorrowingLiteral()
Dxla_compiler.cc453 xla::Shape xla_shape; in XlaCompiler() local
454 TF_RETURN_IF_ERROR(TensorShapeToXLAShape(dtype, shape, &xla_shape)); in XlaCompiler()
455 return xla_shape; in XlaCompiler()
614 xla::Shape* xla_shape) const { in XLAShapeForArgument()
627 TF_ASSIGN_OR_RETURN(*xla_shape, in XLAShapeForArgument()
631 *xla_shape = absl::get<xla::Shape>(arg.shape); in XLAShapeForArgument()
634 arg.type, absl::get<TensorShape>(arg.shape), xla_shape)); in XLAShapeForArgument()
645 TF_ASSIGN_OR_RETURN(*xla_shape, in XLAShapeForArgument()
660 TF_RETURN_IF_ERROR(TensorShapeToXLAShape(arg.type, shape, xla_shape)); in XLAShapeForArgument()
664 arg.tensor_array_gradients.size() + 1, *xla_shape); in XLAShapeForArgument()
[all …]
Dxla_op_kernel.cc393 xla::Shape xla_shape; in ReadVariableInputTensor() local
395 TensorShapeToXLAShape(variable->type(), variable->shape(), &xla_shape)); in ReadVariableInputTensor()
396 if (xla::ShapeUtil::Compatible(xla_shape, representation_shape)) { in ReadVariableInputTensor()
531 xla::Shape xla_shape; in AssignVariableTensor() local
532 TF_RETURN_IF_ERROR(TensorShapeToXLAShape(type, shape, &xla_shape)); in AssignVariableTensor()
533 if (!xla::ShapeUtil::Compatible(xla_shape, representation_shape)) { in AssignVariableTensor()
Dxla_expression.cc146 TF_ASSIGN_OR_RETURN(xla::Shape xla_shape, in GetShape()
149 TF_RETURN_IF_ERROR(XLAShapeToTensorShape(xla_shape, &shape)); in GetShape()
Dliteral_util.h46 const xla::Shape& xla_shape, Tensor* host_tensor,
Dxla_compiler_test.cc313 xla::Shape xla_shape; in TEST_F() local
314 TF_RETURN_IF_ERROR(TensorShapeToXLAShape(dt, shape, &xla_shape)); in TEST_F()
315 *xla_shape.mutable_layout() = xla::LayoutUtil::MakeLayout({0, 1}); in TEST_F()
316 return xla_shape; in TEST_F()
1036 xla::Shape xla_shape; in TEST_F() local
1037 TF_RETURN_IF_ERROR(TensorShapeToXLAShape(type, shape, &xla_shape)); in TEST_F()
1038 *xla_shape.mutable_layout() = xla::LayoutUtil::MakeLayout({0, 1}); in TEST_F()
1039 return xla_shape; in TEST_F()
1074 xla::Shape xla_shape; in TEST_F() local
1075 TF_RETURN_IF_ERROR(TensorShapeToXLAShape(type, shape, &xla_shape)); in TEST_F()
[all …]
Dxla_compiler.h360 xla::Shape* xla_shape) const;
/external/tensorflow/tensorflow/compiler/xla/python_api/
DBUILD20 name = "xla_shape",
21 srcs = ["xla_shape.py"],
35 ":xla_shape",
Dxla_literal.py25 from tensorflow.compiler.xla.python_api import xla_shape
68 literal.shape.CopyFrom(xla_shape.CreateShapeFromNumpy(ndarray).message)
89 literal.shape.CopyFrom(xla_shape.CreateShapeFromNumpy(value).message)
/external/tensorflow/tensorflow/compiler/jit/
Dxla_device_context.cc92 xla::Shape xla_shape; in XlaDeviceContext() local
93 TF_RETURN_IF_ERROR(TensorShapeToXLAShape(dtype, shape, &xla_shape)); in XlaDeviceContext()
94 return xla_shape; in XlaDeviceContext()
/external/tensorflow/tensorflow/compiler/xla/experimental/xla_sharding/
DBUILD15 "//tensorflow/compiler/xla/python_api:xla_shape",
/external/tensorflow/tensorflow/compiler/xrt/kernels/
Dxrt_state_ops.h230 xla::Shape xla_shape; in XRTAllocateFromTensorOp() local
232 ctx, TensorShapeToXLAShape(dtypes_[i], tf_shapes_[i], &xla_shape)); in XRTAllocateFromTensorOp()
233 xla_shapes.push_back(std::move(xla_shape)); in XRTAllocateFromTensorOp()
/external/tensorflow/tensorflow/python/tpu/
DBUILD147 "//tensorflow/compiler/xla/python_api:xla_shape",
/external/tensorflow/tensorflow/compiler/xla/python/
Dxla_client.py1089 xla_shape = Shape.array_shape(dtype, shape)
1090 return self._client.BroadcastedIota(xla_shape, dimension)
Dxla_client_test.py503 xla_shape = local_buffer.shape()
504 self.assertEqual(xla_shape.dimensions(), (1, 2,))
505 self.assertEqual(np.dtype(xla_shape.element_type()), np.dtype(np.float32))