376 lines
16 KiB
Python
376 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
|
||
|
|
|
||
|
|
def tpu_partitioned_input(inputs, partition_dim=0, name=None):
|
||
|
|
r"""An op that groups a list of partitioned inputs together. This op
|
||
|
|
|
||
|
|
Args:
|
||
|
|
inputs: A list of at least 1 `Tensor` objects with the same type.
|
||
|
|
A list of partitioned inputs which must have the same shape.
|
||
|
|
partition_dim: An optional `int`. Defaults to `0`.
|
||
|
|
An integer describles which dimension is partitioned. -1 means
|
||
|
|
those inputs are replicated.
|
||
|
|
name: A name for the operation (optional).
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
A `Tensor`. Has the same type as `inputs`.
|
||
|
|
"""
|
||
|
|
_ctx = _context._context or _context.context()
|
||
|
|
tld = _ctx._thread_local_data
|
||
|
|
if tld.is_eager:
|
||
|
|
try:
|
||
|
|
_result = pywrap_tfe.TFE_Py_FastPathExecute(
|
||
|
|
_ctx, "TPUPartitionedInput", name, inputs, "partition_dim",
|
||
|
|
partition_dim)
|
||
|
|
return _result
|
||
|
|
except _core._NotOkStatusException as e:
|
||
|
|
_ops.raise_from_not_ok_status(e, name)
|
||
|
|
except _core._FallbackException:
|
||
|
|
pass
|
||
|
|
try:
|
||
|
|
return tpu_partitioned_input_eager_fallback(
|
||
|
|
inputs, partition_dim=partition_dim, name=name, ctx=_ctx)
|
||
|
|
except _core._SymbolicException:
|
||
|
|
pass # Add nodes to the TensorFlow graph.
|
||
|
|
# Add nodes to the TensorFlow graph.
|
||
|
|
if not isinstance(inputs, (list, tuple)):
|
||
|
|
raise TypeError(
|
||
|
|
"Expected list for 'inputs' argument to "
|
||
|
|
"'tpu_partitioned_input' Op, not %r." % inputs)
|
||
|
|
_attr_N = len(inputs)
|
||
|
|
if partition_dim is None:
|
||
|
|
partition_dim = 0
|
||
|
|
partition_dim = _execute.make_int(partition_dim, "partition_dim")
|
||
|
|
_, _, _op, _outputs = _op_def_library._apply_op_helper(
|
||
|
|
"TPUPartitionedInput", inputs=inputs, partition_dim=partition_dim,
|
||
|
|
name=name)
|
||
|
|
_result = _outputs[:]
|
||
|
|
if _execute.must_record_gradient():
|
||
|
|
_attrs = ("N", _op._get_attr_int("N"), "T", _op._get_attr_type("T"),
|
||
|
|
"partition_dim", _op._get_attr_int("partition_dim"))
|
||
|
|
_inputs_flat = _op.inputs
|
||
|
|
_execute.record_gradient(
|
||
|
|
"TPUPartitionedInput", _inputs_flat, _attrs, _result)
|
||
|
|
_result, = _result
|
||
|
|
return _result
|
||
|
|
|
||
|
|
TPUPartitionedInput = tf_export("raw_ops.TPUPartitionedInput")(_ops.to_raw_op(tpu_partitioned_input))
|
||
|
|
|
||
|
|
|
||
|
|
def tpu_partitioned_input_eager_fallback(inputs, partition_dim, name, ctx):
|
||
|
|
if not isinstance(inputs, (list, tuple)):
|
||
|
|
raise TypeError(
|
||
|
|
"Expected list for 'inputs' argument to "
|
||
|
|
"'tpu_partitioned_input' Op, not %r." % inputs)
|
||
|
|
_attr_N = len(inputs)
|
||
|
|
if partition_dim is None:
|
||
|
|
partition_dim = 0
|
||
|
|
partition_dim = _execute.make_int(partition_dim, "partition_dim")
|
||
|
|
_attr_T, inputs = _execute.args_to_matching_eager(list(inputs), ctx, [])
|
||
|
|
_inputs_flat = list(inputs)
|
||
|
|
_attrs = ("N", _attr_N, "T", _attr_T, "partition_dim", partition_dim)
|
||
|
|
_result = _execute.execute(b"TPUPartitionedInput", 1, inputs=_inputs_flat,
|
||
|
|
attrs=_attrs, ctx=ctx, name=name)
|
||
|
|
if _execute.must_record_gradient():
|
||
|
|
_execute.record_gradient(
|
||
|
|
"TPUPartitionedInput", _inputs_flat, _attrs, _result)
|
||
|
|
_result, = _result
|
||
|
|
return _result
|
||
|
|
|
||
|
|
|
||
|
|
def tpu_partitioned_input_v2(inputs, partition_dims, is_packed=False, name=None):
|
||
|
|
r"""An op that groups a list of partitioned inputs together. Supports ND sharding.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
inputs: A list of at least 1 `Tensor` objects with the same type.
|
||
|
|
A list of partitioned inputs which must have the same shape.
|
||
|
|
partition_dims: A list of `ints`.
|
||
|
|
A list of integers describing how each dimension is partitioned. Emptiness
|
||
|
|
indicates the inputs are replicated.
|
||
|
|
is_packed: An optional `bool`. Defaults to `False`.
|
||
|
|
Indicates whether the input is a packed resource.
|
||
|
|
name: A name for the operation (optional).
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
A `Tensor`. Has the same type as `inputs`.
|
||
|
|
"""
|
||
|
|
_ctx = _context._context or _context.context()
|
||
|
|
tld = _ctx._thread_local_data
|
||
|
|
if tld.is_eager:
|
||
|
|
try:
|
||
|
|
_result = pywrap_tfe.TFE_Py_FastPathExecute(
|
||
|
|
_ctx, "TPUPartitionedInputV2", name, inputs, "partition_dims",
|
||
|
|
partition_dims, "is_packed", is_packed)
|
||
|
|
return _result
|
||
|
|
except _core._NotOkStatusException as e:
|
||
|
|
_ops.raise_from_not_ok_status(e, name)
|
||
|
|
except _core._FallbackException:
|
||
|
|
pass
|
||
|
|
try:
|
||
|
|
return tpu_partitioned_input_v2_eager_fallback(
|
||
|
|
inputs, partition_dims=partition_dims, is_packed=is_packed,
|
||
|
|
name=name, ctx=_ctx)
|
||
|
|
except _core._SymbolicException:
|
||
|
|
pass # Add nodes to the TensorFlow graph.
|
||
|
|
# Add nodes to the TensorFlow graph.
|
||
|
|
if not isinstance(inputs, (list, tuple)):
|
||
|
|
raise TypeError(
|
||
|
|
"Expected list for 'inputs' argument to "
|
||
|
|
"'tpu_partitioned_input_v2' Op, not %r." % inputs)
|
||
|
|
_attr_N = len(inputs)
|
||
|
|
if not isinstance(partition_dims, (list, tuple)):
|
||
|
|
raise TypeError(
|
||
|
|
"Expected list for 'partition_dims' argument to "
|
||
|
|
"'tpu_partitioned_input_v2' Op, not %r." % partition_dims)
|
||
|
|
partition_dims = [_execute.make_int(_i, "partition_dims") for _i in partition_dims]
|
||
|
|
if is_packed is None:
|
||
|
|
is_packed = False
|
||
|
|
is_packed = _execute.make_bool(is_packed, "is_packed")
|
||
|
|
_, _, _op, _outputs = _op_def_library._apply_op_helper(
|
||
|
|
"TPUPartitionedInputV2", inputs=inputs, partition_dims=partition_dims,
|
||
|
|
is_packed=is_packed, name=name)
|
||
|
|
_result = _outputs[:]
|
||
|
|
if _execute.must_record_gradient():
|
||
|
|
_attrs = ("N", _op._get_attr_int("N"), "T", _op._get_attr_type("T"),
|
||
|
|
"partition_dims", _op.get_attr("partition_dims"), "is_packed",
|
||
|
|
_op._get_attr_bool("is_packed"))
|
||
|
|
_inputs_flat = _op.inputs
|
||
|
|
_execute.record_gradient(
|
||
|
|
"TPUPartitionedInputV2", _inputs_flat, _attrs, _result)
|
||
|
|
_result, = _result
|
||
|
|
return _result
|
||
|
|
|
||
|
|
TPUPartitionedInputV2 = tf_export("raw_ops.TPUPartitionedInputV2")(_ops.to_raw_op(tpu_partitioned_input_v2))
|
||
|
|
|
||
|
|
|
||
|
|
def tpu_partitioned_input_v2_eager_fallback(inputs, partition_dims, is_packed, name, ctx):
|
||
|
|
if not isinstance(inputs, (list, tuple)):
|
||
|
|
raise TypeError(
|
||
|
|
"Expected list for 'inputs' argument to "
|
||
|
|
"'tpu_partitioned_input_v2' Op, not %r." % inputs)
|
||
|
|
_attr_N = len(inputs)
|
||
|
|
if not isinstance(partition_dims, (list, tuple)):
|
||
|
|
raise TypeError(
|
||
|
|
"Expected list for 'partition_dims' argument to "
|
||
|
|
"'tpu_partitioned_input_v2' Op, not %r." % partition_dims)
|
||
|
|
partition_dims = [_execute.make_int(_i, "partition_dims") for _i in partition_dims]
|
||
|
|
if is_packed is None:
|
||
|
|
is_packed = False
|
||
|
|
is_packed = _execute.make_bool(is_packed, "is_packed")
|
||
|
|
_attr_T, inputs = _execute.args_to_matching_eager(list(inputs), ctx, [])
|
||
|
|
_inputs_flat = list(inputs)
|
||
|
|
_attrs = ("N", _attr_N, "T", _attr_T, "partition_dims", partition_dims,
|
||
|
|
"is_packed", is_packed)
|
||
|
|
_result = _execute.execute(b"TPUPartitionedInputV2", 1, inputs=_inputs_flat,
|
||
|
|
attrs=_attrs, ctx=ctx, name=name)
|
||
|
|
if _execute.must_record_gradient():
|
||
|
|
_execute.record_gradient(
|
||
|
|
"TPUPartitionedInputV2", _inputs_flat, _attrs, _result)
|
||
|
|
_result, = _result
|
||
|
|
return _result
|
||
|
|
|
||
|
|
|
||
|
|
def tpu_partitioned_output(inputs, num_splits, partition_dim=0, name=None):
|
||
|
|
r"""An op that demultiplexes a tensor to be sharded by XLA to a list of partitioned
|
||
|
|
|
||
|
|
outputs outside the XLA computation.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
inputs: A `Tensor`.
|
||
|
|
A tensor which represents the full shape of partitioned tensors.
|
||
|
|
num_splits: An `int` that is `>= 1`.
|
||
|
|
partition_dim: An optional `int`. Defaults to `0`.
|
||
|
|
An integer describles which dimension is partitioned.
|
||
|
|
name: A name for the operation (optional).
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
A list of `num_splits` `Tensor` objects with the same type as `inputs`.
|
||
|
|
"""
|
||
|
|
_ctx = _context._context or _context.context()
|
||
|
|
tld = _ctx._thread_local_data
|
||
|
|
if tld.is_eager:
|
||
|
|
try:
|
||
|
|
_result = pywrap_tfe.TFE_Py_FastPathExecute(
|
||
|
|
_ctx, "TPUPartitionedOutput", name, inputs, "num_splits", num_splits,
|
||
|
|
"partition_dim", partition_dim)
|
||
|
|
return _result
|
||
|
|
except _core._NotOkStatusException as e:
|
||
|
|
_ops.raise_from_not_ok_status(e, name)
|
||
|
|
except _core._FallbackException:
|
||
|
|
pass
|
||
|
|
try:
|
||
|
|
return tpu_partitioned_output_eager_fallback(
|
||
|
|
inputs, num_splits=num_splits, partition_dim=partition_dim,
|
||
|
|
name=name, ctx=_ctx)
|
||
|
|
except _core._SymbolicException:
|
||
|
|
pass # Add nodes to the TensorFlow graph.
|
||
|
|
# Add nodes to the TensorFlow graph.
|
||
|
|
num_splits = _execute.make_int(num_splits, "num_splits")
|
||
|
|
if partition_dim is None:
|
||
|
|
partition_dim = 0
|
||
|
|
partition_dim = _execute.make_int(partition_dim, "partition_dim")
|
||
|
|
_, _, _op, _outputs = _op_def_library._apply_op_helper(
|
||
|
|
"TPUPartitionedOutput", inputs=inputs, num_splits=num_splits,
|
||
|
|
partition_dim=partition_dim, name=name)
|
||
|
|
_result = _outputs[:]
|
||
|
|
if _execute.must_record_gradient():
|
||
|
|
_attrs = ("T", _op._get_attr_type("T"), "num_splits",
|
||
|
|
_op._get_attr_int("num_splits"), "partition_dim",
|
||
|
|
_op._get_attr_int("partition_dim"))
|
||
|
|
_inputs_flat = _op.inputs
|
||
|
|
_execute.record_gradient(
|
||
|
|
"TPUPartitionedOutput", _inputs_flat, _attrs, _result)
|
||
|
|
return _result
|
||
|
|
|
||
|
|
TPUPartitionedOutput = tf_export("raw_ops.TPUPartitionedOutput")(_ops.to_raw_op(tpu_partitioned_output))
|
||
|
|
|
||
|
|
|
||
|
|
def tpu_partitioned_output_eager_fallback(inputs, num_splits, partition_dim, name, ctx):
|
||
|
|
num_splits = _execute.make_int(num_splits, "num_splits")
|
||
|
|
if partition_dim is None:
|
||
|
|
partition_dim = 0
|
||
|
|
partition_dim = _execute.make_int(partition_dim, "partition_dim")
|
||
|
|
_attr_T, (inputs,) = _execute.args_to_matching_eager([inputs], ctx, [])
|
||
|
|
_inputs_flat = [inputs]
|
||
|
|
_attrs = ("T", _attr_T, "num_splits", num_splits, "partition_dim",
|
||
|
|
partition_dim)
|
||
|
|
_result = _execute.execute(b"TPUPartitionedOutput", num_splits,
|
||
|
|
inputs=_inputs_flat, attrs=_attrs, ctx=ctx,
|
||
|
|
name=name)
|
||
|
|
if _execute.must_record_gradient():
|
||
|
|
_execute.record_gradient(
|
||
|
|
"TPUPartitionedOutput", _inputs_flat, _attrs, _result)
|
||
|
|
return _result
|
||
|
|
|
||
|
|
|
||
|
|
@_dispatch.add_fallback_dispatch_list
|
||
|
|
@_dispatch.add_type_based_api_dispatcher
|
||
|
|
@tf_export('tpu_partitioned_output_v2')
|
||
|
|
def tpu_partitioned_output_v2(inputs, num_splits, partition_dims, name=None):
|
||
|
|
r"""An op that demultiplexes a tensor to be sharded by XLA to a list of partitioned
|
||
|
|
|
||
|
|
outputs outside the XLA computation. Supports ND sharding.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
inputs: A `Tensor`.
|
||
|
|
A tensor which represents the full shape of partitioned tensors.
|
||
|
|
num_splits: An `int` that is `>= 1`.
|
||
|
|
partition_dims: A list of `ints`.
|
||
|
|
A list of integers describing how each dimension is partitioned. Emptiness
|
||
|
|
indicates the inputs are replicated.
|
||
|
|
name: A name for the operation (optional).
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
A list of `num_splits` `Tensor` objects with the same type as `inputs`.
|
||
|
|
"""
|
||
|
|
_ctx = _context._context or _context.context()
|
||
|
|
tld = _ctx._thread_local_data
|
||
|
|
if tld.is_eager:
|
||
|
|
try:
|
||
|
|
_result = pywrap_tfe.TFE_Py_FastPathExecute(
|
||
|
|
_ctx, "TPUPartitionedOutputV2", name, inputs, "num_splits",
|
||
|
|
num_splits, "partition_dims", partition_dims)
|
||
|
|
return _result
|
||
|
|
except _core._NotOkStatusException as e:
|
||
|
|
_ops.raise_from_not_ok_status(e, name)
|
||
|
|
except _core._FallbackException:
|
||
|
|
pass
|
||
|
|
try:
|
||
|
|
_result = _dispatcher_for_tpu_partitioned_output_v2(
|
||
|
|
(inputs, num_splits, partition_dims, name,), None)
|
||
|
|
if _result is not NotImplemented:
|
||
|
|
return _result
|
||
|
|
return tpu_partitioned_output_v2_eager_fallback(
|
||
|
|
inputs, num_splits=num_splits, partition_dims=partition_dims,
|
||
|
|
name=name, ctx=_ctx)
|
||
|
|
except _core._SymbolicException:
|
||
|
|
pass # Add nodes to the TensorFlow graph.
|
||
|
|
except (TypeError, ValueError):
|
||
|
|
_result = _dispatch.dispatch(
|
||
|
|
tpu_partitioned_output_v2, (), dict(inputs=inputs,
|
||
|
|
num_splits=num_splits,
|
||
|
|
partition_dims=partition_dims,
|
||
|
|
name=name)
|
||
|
|
)
|
||
|
|
if _result is not _dispatch.OpDispatcher.NOT_SUPPORTED:
|
||
|
|
return _result
|
||
|
|
raise
|
||
|
|
else:
|
||
|
|
_result = _dispatcher_for_tpu_partitioned_output_v2(
|
||
|
|
(inputs, num_splits, partition_dims, name,), None)
|
||
|
|
if _result is not NotImplemented:
|
||
|
|
return _result
|
||
|
|
# Add nodes to the TensorFlow graph.
|
||
|
|
num_splits = _execute.make_int(num_splits, "num_splits")
|
||
|
|
if not isinstance(partition_dims, (list, tuple)):
|
||
|
|
raise TypeError(
|
||
|
|
"Expected list for 'partition_dims' argument to "
|
||
|
|
"'tpu_partitioned_output_v2' Op, not %r." % partition_dims)
|
||
|
|
partition_dims = [_execute.make_int(_i, "partition_dims") for _i in partition_dims]
|
||
|
|
try:
|
||
|
|
_, _, _op, _outputs = _op_def_library._apply_op_helper(
|
||
|
|
"TPUPartitionedOutputV2", inputs=inputs, num_splits=num_splits,
|
||
|
|
partition_dims=partition_dims, name=name)
|
||
|
|
except (TypeError, ValueError):
|
||
|
|
_result = _dispatch.dispatch(
|
||
|
|
tpu_partitioned_output_v2, (), dict(inputs=inputs,
|
||
|
|
num_splits=num_splits,
|
||
|
|
partition_dims=partition_dims,
|
||
|
|
name=name)
|
||
|
|
)
|
||
|
|
if _result is not _dispatch.OpDispatcher.NOT_SUPPORTED:
|
||
|
|
return _result
|
||
|
|
raise
|
||
|
|
_result = _outputs[:]
|
||
|
|
if _execute.must_record_gradient():
|
||
|
|
_attrs = ("T", _op._get_attr_type("T"), "num_splits",
|
||
|
|
_op._get_attr_int("num_splits"), "partition_dims",
|
||
|
|
_op.get_attr("partition_dims"))
|
||
|
|
_inputs_flat = _op.inputs
|
||
|
|
_execute.record_gradient(
|
||
|
|
"TPUPartitionedOutputV2", _inputs_flat, _attrs, _result)
|
||
|
|
return _result
|
||
|
|
|
||
|
|
TPUPartitionedOutputV2 = tf_export("raw_ops.TPUPartitionedOutputV2")(_ops.to_raw_op(tpu_partitioned_output_v2))
|
||
|
|
_dispatcher_for_tpu_partitioned_output_v2 = tpu_partitioned_output_v2._tf_type_based_dispatcher.Dispatch
|
||
|
|
|
||
|
|
|
||
|
|
def tpu_partitioned_output_v2_eager_fallback(inputs, num_splits, partition_dims, name, ctx):
|
||
|
|
num_splits = _execute.make_int(num_splits, "num_splits")
|
||
|
|
if not isinstance(partition_dims, (list, tuple)):
|
||
|
|
raise TypeError(
|
||
|
|
"Expected list for 'partition_dims' argument to "
|
||
|
|
"'tpu_partitioned_output_v2' Op, not %r." % partition_dims)
|
||
|
|
partition_dims = [_execute.make_int(_i, "partition_dims") for _i in partition_dims]
|
||
|
|
_attr_T, (inputs,) = _execute.args_to_matching_eager([inputs], ctx, [])
|
||
|
|
_inputs_flat = [inputs]
|
||
|
|
_attrs = ("T", _attr_T, "num_splits", num_splits, "partition_dims",
|
||
|
|
partition_dims)
|
||
|
|
_result = _execute.execute(b"TPUPartitionedOutputV2", num_splits,
|
||
|
|
inputs=_inputs_flat, attrs=_attrs, ctx=ctx,
|
||
|
|
name=name)
|
||
|
|
if _execute.must_record_gradient():
|
||
|
|
_execute.record_gradient(
|
||
|
|
"TPUPartitionedOutputV2", _inputs_flat, _attrs, _result)
|
||
|
|
return _result
|
||
|
|
|