Returns the size of a tensor. Args: input: A `Tensor` or `SparseTensor`. name: A name for the operation (optional). optimize: if true, encode the size as a constant when possible. out_type: (Optional) The specified non-quantized numeric output type of the operation. Defaults
(input, name=None, optimize=True, out_type=dtypes.int32)
| 534 | |
| 535 | |
| 536 | def size_internal(input, name=None, optimize=True, out_type=dtypes.int32): |
| 537 | # pylint: disable=redefined-builtin,protected-access |
| 538 | """Returns the size of a tensor. |
| 539 | |
| 540 | Args: |
| 541 | input: A `Tensor` or `SparseTensor`. |
| 542 | name: A name for the operation (optional). |
| 543 | optimize: if true, encode the size as a constant when possible. |
| 544 | out_type: (Optional) The specified non-quantized numeric output type of the |
| 545 | operation. Defaults to `tf.int32`. |
| 546 | |
| 547 | Returns: |
| 548 | A `Tensor` of type `out_type`. Defaults to `tf.int32`. |
| 549 | """ |
| 550 | if (context.executing_eagerly() |
| 551 | and not hasattr(input, "graph") |
| 552 | and not isinstance( |
| 553 | input, (sparse_tensor.SparseTensor, sparse_tensor.SparseTensorValue)) |
| 554 | ): |
| 555 | input = ops.convert_to_tensor(input) |
| 556 | np_out_type = out_type.as_numpy_dtype |
| 557 | num_elements = np.prod(input._shape_tuple(), dtype=np_out_type) # pylint: disable=protected-access |
| 558 | return ops.convert_to_tensor(num_elements, dtype=out_type) |
| 559 | with ops.name_scope(name, "Size", [input]) as name: |
| 560 | if isinstance( |
| 561 | input, (sparse_tensor.SparseTensor, sparse_tensor.SparseTensorValue)): |
| 562 | return gen_math_ops.prod( |
| 563 | gen_math_ops.cast(input.dense_shape, out_type), 0, name=name) |
| 564 | else: |
| 565 | input_tensor = ops.convert_to_tensor(input) |
| 566 | input_shape = input_tensor.get_shape() |
| 567 | if optimize: |
| 568 | if input_shape.is_fully_defined(): |
| 569 | return constant(input_shape.num_elements(), out_type, name=name) |
| 570 | if input_shape.dims and any(dim == 0 for dim in input_shape.dims): |
| 571 | return constant(0, out_type, name=name) |
| 572 | return gen_array_ops.size(input, name=name, out_type=out_type) |
| 573 | |
| 574 | |
| 575 | @tf_export("rank") |
no test coverage detected