308 lines
16 KiB
Python
308 lines
16 KiB
Python
"""Python wrappers around TensorFlow ops.
|
|
|
|
This file is MACHINE GENERATED! Do not edit.
|
|
"""
|
|
|
|
import collections
|
|
|
|
from tensorflow.python import pywrap_tfe as pywrap_tfe
|
|
from tensorflow.python.eager import context as _context
|
|
from tensorflow.python.eager import core as _core
|
|
from tensorflow.python.eager import execute as _execute
|
|
from tensorflow.python.framework import dtypes as _dtypes
|
|
|
|
from tensorflow.python.framework import op_def_registry as _op_def_registry
|
|
from tensorflow.python.framework import ops as _ops
|
|
from tensorflow.python.framework import op_def_library as _op_def_library
|
|
from tensorflow.python.util.deprecation import deprecated_endpoints
|
|
from tensorflow.python.util import dispatch as _dispatch
|
|
from tensorflow.python.util.tf_export import tf_export
|
|
|
|
from typing import TypeVar
|
|
_RaggedCrossOutput = collections.namedtuple(
|
|
"RaggedCross",
|
|
["output_values", "output_row_splits"])
|
|
|
|
|
|
def ragged_cross(ragged_values, ragged_row_splits, sparse_indices, sparse_values, sparse_shape, dense_inputs, input_order, hashed_output, num_buckets, hash_key, out_values_type, out_row_splits_type, name=None):
|
|
r"""Generates a feature cross from a list of tensors, and returns it as a
|
|
|
|
RaggedTensor. See `tf.ragged.cross` for more details.
|
|
|
|
|
|
Args:
|
|
ragged_values: A list of `Tensor` objects with types from: `int64`, `string`.
|
|
The values tensor for each RaggedTensor input.
|
|
ragged_row_splits: A list of `Tensor` objects with types from: `int32`, `int64`.
|
|
The row_splits tensor for each RaggedTensor input.
|
|
sparse_indices: A list of `Tensor` objects with type `int64`.
|
|
The indices tensor for each SparseTensor input.
|
|
sparse_values: A list of `Tensor` objects with types from: `int64`, `string`.
|
|
The values tensor for each SparseTensor input.
|
|
sparse_shape: A list with the same length as `sparse_indices` of `Tensor` objects with type `int64`.
|
|
The dense_shape tensor for each SparseTensor input.
|
|
dense_inputs: A list of `Tensor` objects with types from: `int64`, `string`.
|
|
The tf.Tensor inputs.
|
|
input_order: A `string`.
|
|
String specifying the tensor type for each input. The `i`th character in
|
|
|
|
this string specifies the type of the `i`th input, and is one of: 'R' (ragged),
|
|
|
|
'D' (dense), or 'S' (sparse). This attr is used to ensure that the crossed
|
|
|
|
values are combined in the order of the inputs from the call to tf.ragged.cross.
|
|
hashed_output: A `bool`.
|
|
num_buckets: An `int` that is `>= 0`.
|
|
hash_key: An `int`.
|
|
out_values_type: A `tf.DType` from: `tf.int64, tf.string`.
|
|
out_row_splits_type: A `tf.DType` from: `tf.int32, tf.int64`.
|
|
name: A name for the operation (optional).
|
|
|
|
Returns:
|
|
A tuple of `Tensor` objects (output_values, output_row_splits).
|
|
|
|
output_values: A `Tensor` of type `out_values_type`.
|
|
output_row_splits: A `Tensor` of type `out_row_splits_type`.
|
|
"""
|
|
_ctx = _context._context or _context.context()
|
|
tld = _ctx._thread_local_data
|
|
if tld.is_eager:
|
|
try:
|
|
_result = pywrap_tfe.TFE_Py_FastPathExecute(
|
|
_ctx, "RaggedCross", name, ragged_values, ragged_row_splits,
|
|
sparse_indices, sparse_values, sparse_shape, dense_inputs,
|
|
"input_order", input_order, "hashed_output", hashed_output,
|
|
"num_buckets", num_buckets, "hash_key", hash_key, "out_values_type",
|
|
out_values_type, "out_row_splits_type", out_row_splits_type)
|
|
_result = _RaggedCrossOutput._make(_result)
|
|
return _result
|
|
except _core._NotOkStatusException as e:
|
|
_ops.raise_from_not_ok_status(e, name)
|
|
except _core._FallbackException:
|
|
pass
|
|
try:
|
|
return ragged_cross_eager_fallback(
|
|
ragged_values, ragged_row_splits, sparse_indices, sparse_values,
|
|
sparse_shape, dense_inputs, input_order=input_order,
|
|
hashed_output=hashed_output, num_buckets=num_buckets,
|
|
hash_key=hash_key, out_values_type=out_values_type,
|
|
out_row_splits_type=out_row_splits_type, name=name, ctx=_ctx)
|
|
except _core._SymbolicException:
|
|
pass # Add nodes to the TensorFlow graph.
|
|
# Add nodes to the TensorFlow graph.
|
|
if not isinstance(sparse_indices, (list, tuple)):
|
|
raise TypeError(
|
|
"Expected list for 'sparse_indices' argument to "
|
|
"'ragged_cross' Op, not %r." % sparse_indices)
|
|
_attr_Nsparse = len(sparse_indices)
|
|
if not isinstance(sparse_shape, (list, tuple)):
|
|
raise TypeError(
|
|
"Expected list for 'sparse_shape' argument to "
|
|
"'ragged_cross' Op, not %r." % sparse_shape)
|
|
if len(sparse_shape) != _attr_Nsparse:
|
|
raise ValueError(
|
|
"List argument 'sparse_shape' to 'ragged_cross' Op with length %d "
|
|
"must match length %d of argument 'sparse_indices'." %
|
|
(len(sparse_shape), _attr_Nsparse))
|
|
input_order = _execute.make_str(input_order, "input_order")
|
|
hashed_output = _execute.make_bool(hashed_output, "hashed_output")
|
|
num_buckets = _execute.make_int(num_buckets, "num_buckets")
|
|
hash_key = _execute.make_int(hash_key, "hash_key")
|
|
out_values_type = _execute.make_type(out_values_type, "out_values_type")
|
|
out_row_splits_type = _execute.make_type(out_row_splits_type, "out_row_splits_type")
|
|
_, _, _op, _outputs = _op_def_library._apply_op_helper(
|
|
"RaggedCross", ragged_values=ragged_values,
|
|
ragged_row_splits=ragged_row_splits,
|
|
sparse_indices=sparse_indices,
|
|
sparse_values=sparse_values, sparse_shape=sparse_shape,
|
|
dense_inputs=dense_inputs, input_order=input_order,
|
|
hashed_output=hashed_output, num_buckets=num_buckets,
|
|
hash_key=hash_key, out_values_type=out_values_type,
|
|
out_row_splits_type=out_row_splits_type, name=name)
|
|
_result = _outputs[:]
|
|
if _execute.must_record_gradient():
|
|
_attrs = ("Nsparse", _op._get_attr_int("Nsparse"), "input_order",
|
|
_op.get_attr("input_order"), "hashed_output",
|
|
_op._get_attr_bool("hashed_output"), "num_buckets",
|
|
_op._get_attr_int("num_buckets"), "hash_key",
|
|
_op._get_attr_int("hash_key"), "ragged_values_types",
|
|
_op.get_attr("ragged_values_types"), "ragged_splits_types",
|
|
_op.get_attr("ragged_splits_types"), "sparse_values_types",
|
|
_op.get_attr("sparse_values_types"), "dense_types",
|
|
_op.get_attr("dense_types"), "out_values_type",
|
|
_op._get_attr_type("out_values_type"), "out_row_splits_type",
|
|
_op._get_attr_type("out_row_splits_type"))
|
|
_inputs_flat = _op.inputs
|
|
_execute.record_gradient(
|
|
"RaggedCross", _inputs_flat, _attrs, _result)
|
|
_result = _RaggedCrossOutput._make(_result)
|
|
return _result
|
|
|
|
RaggedCross = tf_export("raw_ops.RaggedCross")(_ops.to_raw_op(ragged_cross))
|
|
|
|
|
|
def ragged_cross_eager_fallback(ragged_values, ragged_row_splits, sparse_indices, sparse_values, sparse_shape, dense_inputs, input_order, hashed_output, num_buckets, hash_key, out_values_type, out_row_splits_type, name, ctx):
|
|
if not isinstance(sparse_indices, (list, tuple)):
|
|
raise TypeError(
|
|
"Expected list for 'sparse_indices' argument to "
|
|
"'ragged_cross' Op, not %r." % sparse_indices)
|
|
_attr_Nsparse = len(sparse_indices)
|
|
if not isinstance(sparse_shape, (list, tuple)):
|
|
raise TypeError(
|
|
"Expected list for 'sparse_shape' argument to "
|
|
"'ragged_cross' Op, not %r." % sparse_shape)
|
|
if len(sparse_shape) != _attr_Nsparse:
|
|
raise ValueError(
|
|
"List argument 'sparse_shape' to 'ragged_cross' Op with length %d "
|
|
"must match length %d of argument 'sparse_indices'." %
|
|
(len(sparse_shape), _attr_Nsparse))
|
|
input_order = _execute.make_str(input_order, "input_order")
|
|
hashed_output = _execute.make_bool(hashed_output, "hashed_output")
|
|
num_buckets = _execute.make_int(num_buckets, "num_buckets")
|
|
hash_key = _execute.make_int(hash_key, "hash_key")
|
|
out_values_type = _execute.make_type(out_values_type, "out_values_type")
|
|
out_row_splits_type = _execute.make_type(out_row_splits_type, "out_row_splits_type")
|
|
_attr_ragged_values_types, ragged_values = _execute.convert_to_mixed_eager_tensors(ragged_values, ctx)
|
|
_attr_ragged_splits_types, ragged_row_splits = _execute.convert_to_mixed_eager_tensors(ragged_row_splits, ctx)
|
|
_attr_sparse_values_types, sparse_values = _execute.convert_to_mixed_eager_tensors(sparse_values, ctx)
|
|
_attr_dense_types, dense_inputs = _execute.convert_to_mixed_eager_tensors(dense_inputs, ctx)
|
|
sparse_indices = _ops.convert_n_to_tensor(sparse_indices, _dtypes.int64)
|
|
sparse_shape = _ops.convert_n_to_tensor(sparse_shape, _dtypes.int64)
|
|
_inputs_flat = list(ragged_values) + list(ragged_row_splits) + list(sparse_indices) + list(sparse_values) + list(sparse_shape) + list(dense_inputs)
|
|
_attrs = ("Nsparse", _attr_Nsparse, "input_order", input_order,
|
|
"hashed_output", hashed_output, "num_buckets", num_buckets, "hash_key",
|
|
hash_key, "ragged_values_types", _attr_ragged_values_types,
|
|
"ragged_splits_types", _attr_ragged_splits_types, "sparse_values_types",
|
|
_attr_sparse_values_types, "dense_types", _attr_dense_types,
|
|
"out_values_type", out_values_type, "out_row_splits_type",
|
|
out_row_splits_type)
|
|
_result = _execute.execute(b"RaggedCross", 2, inputs=_inputs_flat,
|
|
attrs=_attrs, ctx=ctx, name=name)
|
|
if _execute.must_record_gradient():
|
|
_execute.record_gradient(
|
|
"RaggedCross", _inputs_flat, _attrs, _result)
|
|
_result = _RaggedCrossOutput._make(_result)
|
|
return _result
|
|
|
|
_RaggedGatherOutput = collections.namedtuple(
|
|
"RaggedGather",
|
|
["output_nested_splits", "output_dense_values"])
|
|
|
|
|
|
def ragged_gather(params_nested_splits, params_dense_values, indices, OUTPUT_RAGGED_RANK, name=None):
|
|
r"""Gather ragged slices from `params` axis `0` according to `indices`.
|
|
|
|
|
|
Outputs a `RaggedTensor` output composed from `output_dense_values` and
|
|
|
|
`output_nested_splits`, such that:
|
|
|
|
|
|
|
|
```python
|
|
|
|
output.shape = indices.shape + params.shape[1:]
|
|
|
|
output.ragged_rank = indices.shape.ndims + params.ragged_rank
|
|
|
|
output[i...j, d0...dn] = params[indices[i...j], d0...dn]
|
|
|
|
```
|
|
|
|
|
|
|
|
where
|
|
|
|
|
|
|
|
* `params =
|
|
|
|
ragged.from_nested_row_splits(params_dense_values, params_nested_splits)`
|
|
|
|
provides the values that should be gathered.
|
|
|
|
* `indices` ia a dense tensor with dtype `int32` or `int64`, indicating which
|
|
|
|
values should be gathered.
|
|
|
|
* `output =
|
|
|
|
ragged.from_nested_row_splits(output_dense_values, output_nested_splits)`
|
|
|
|
is the output tensor.
|
|
|
|
|
|
|
|
(Note: This c++ op is used to implement the higher-level python
|
|
|
|
`tf.ragged.gather` op, which also supports ragged indices.)
|
|
|
|
Args:
|
|
params_nested_splits: A list of at least 1 `Tensor` objects with the same type in: `int32`, `int64`.
|
|
The `nested_row_splits` tensors that define the row-partitioning for the
|
|
|
|
`params` RaggedTensor input.
|
|
params_dense_values: A `Tensor`.
|
|
The `flat_values` for the `params` RaggedTensor. There was a terminology change
|
|
|
|
at the python level from dense_values to flat_values, so dense_values is the
|
|
|
|
deprecated name.
|
|
indices: A `Tensor`. Must be one of the following types: `int32`, `int64`.
|
|
Indices in the outermost dimension of `params` of the values that should be
|
|
|
|
gathered.
|
|
OUTPUT_RAGGED_RANK: An `int` that is `>= 0`.
|
|
The ragged rank of the output RaggedTensor. `output_nested_splits` will contain
|
|
|
|
this number of `row_splits` tensors. This value should equal
|
|
|
|
`indices.shape.ndims + params.ragged_rank - 1`.
|
|
name: A name for the operation (optional).
|
|
|
|
Returns:
|
|
A tuple of `Tensor` objects (output_nested_splits, output_dense_values).
|
|
|
|
output_nested_splits: A list of `OUTPUT_RAGGED_RANK` `Tensor` objects with the same type as `params_nested_splits`.
|
|
output_dense_values: A `Tensor`. Has the same type as `params_dense_values`.
|
|
"""
|
|
_ctx = _context._context or _context.context()
|
|
tld = _ctx._thread_local_data
|
|
if tld.is_eager:
|
|
try:
|
|
_result = pywrap_tfe.TFE_Py_FastPathExecute(
|
|
_ctx, "RaggedGather", name, params_nested_splits, params_dense_values,
|
|
indices, "OUTPUT_RAGGED_RANK", OUTPUT_RAGGED_RANK)
|
|
_result = _RaggedGatherOutput._make(_result)
|
|
return _result
|
|
except _core._NotOkStatusException as e:
|
|
_ops.raise_from_not_ok_status(e, name)
|
|
except _core._FallbackException:
|
|
pass
|
|
try:
|
|
return ragged_gather_eager_fallback(
|
|
params_nested_splits, params_dense_values, indices,
|
|
OUTPUT_RAGGED_RANK=OUTPUT_RAGGED_RANK, name=name, ctx=_ctx)
|
|
except _core._SymbolicException:
|
|
pass # Add nodes to the TensorFlow graph.
|
|
# Add nodes to the TensorFlow graph.
|
|
if not isinstance(params_nested_splits, (list, tuple)):
|
|
raise TypeError(
|
|
"Expected list for 'params_nested_splits' argument to "
|
|
"'ragged_gather' Op, not %r." % params_nested_splits)
|
|
_attr_PARAMS_RAGGED_RANK = len(params_nested_splits)
|
|
OUTPUT_RAGGED_RANK = _execute.make_int(OUTPUT_RAGGED_RANK, "OUTPUT_RAGGED_RANK")
|
|
_, _, _op, _outputs = _op_def_library._apply_op_helper(
|
|
"RaggedGather", params_nested_splits=params_nested_splits,
|
|
params_dense_values=params_dense_values,
|
|
indices=indices,
|
|
OUTPUT_RAGGED_RANK=OUTPUT_RAGGED_RANK, name=name)
|
|
_result = _outputs[:]
|
|
if _execute.must_record_gradient():
|
|
_attrs = ("Tvalues", _op._get_attr_type("Tvalues"), "Tindices",
|
|
_op._get_attr_type("Tindices"), "Tsplits",
|
|
_op._get_attr_type("Tsplits"), "PARAMS_RAGGED_RANK",
|
|
_op._get_attr_int("PARAMS_RAGGED_RANK"), "OUTPUT_RAGGED_RANK",
|
|
_op._get_attr_int("OUTPUT_RAGGED_RANK"))
|
|
_inputs_flat = _op.inputs
|
|
_execute.record_gradient(
|
|
"RaggedGather", _inputs_flat, _attrs, _result)
|
|
_result = [_result[:OUTPUT_RAGGED_RANK]] + _result[OUTPUT_RAGGED_RANK:]
|
|
_result = _RaggedGatherOutput._make(_result)
|
|
return _result
|
|
|
|
RaggedGather = tf_export("raw_ops.RaggedGather")(_ops.to_raw_op(ragged_gather))
|
|
|
|
|
|
def ragged_gather_eager_fallback(params_nested_splits, params_dense_values, indices, OUTPUT_RAGGED_RANK, name, ctx):
|
|
if not isinstance(params_nested_splits, (list, tuple)):
|
|
raise TypeError(
|
|
"Expected list for 'params_nested_splits' argument to "
|
|
"'ragged_gather' Op, not %r." % params_nested_splits)
|
|
_attr_PARAMS_RAGGED_RANK = len(params_nested_splits)
|
|
OUTPUT_RAGGED_RANK = _execute.make_int(OUTPUT_RAGGED_RANK, "OUTPUT_RAGGED_RANK")
|
|
_attr_Tvalues, (params_dense_values,) = _execute.args_to_matching_eager([params_dense_values], ctx, [])
|
|
_attr_Tindices, (indices,) = _execute.args_to_matching_eager([indices], ctx, [_dtypes.int32, _dtypes.int64, ])
|
|
_attr_Tsplits, params_nested_splits = _execute.args_to_matching_eager(list(params_nested_splits), ctx, [_dtypes.int32, _dtypes.int64, ], _dtypes.int64)
|
|
_inputs_flat = list(params_nested_splits) + [params_dense_values, indices]
|
|
_attrs = ("Tvalues", _attr_Tvalues, "Tindices", _attr_Tindices, "Tsplits",
|
|
_attr_Tsplits, "PARAMS_RAGGED_RANK", _attr_PARAMS_RAGGED_RANK,
|
|
"OUTPUT_RAGGED_RANK", OUTPUT_RAGGED_RANK)
|
|
_result = _execute.execute(b"RaggedGather", OUTPUT_RAGGED_RANK + 1,
|
|
inputs=_inputs_flat, attrs=_attrs, ctx=ctx,
|
|
name=name)
|
|
if _execute.must_record_gradient():
|
|
_execute.record_gradient(
|
|
"RaggedGather", _inputs_flat, _attrs, _result)
|
|
_result = [_result[:OUTPUT_RAGGED_RANK]] + _result[OUTPUT_RAGGED_RANK:]
|
|
_result = _RaggedGatherOutput._make(_result)
|
|
return _result
|
|
|