TF 2.0 is out! Get hands-on practice at TF World, Oct 28-31. Use code TF20 for 20% off select passes. Register now

tfp.experimental.auto_batching.NumpyBackend

View source on GitHub

Class NumpyBackend

Implements the Numpy backend ops for a PC auto-batching VM.

Aliases:

Properties

variable_class

Methods

any

View source

any(
    t,
    name=None
)

assert_matching_dtype

View source

assert_matching_dtype(
    expected_dtype,
    val,
    message=''
)

Asserts that the dtype of val matches expected_dtype.

Args:

  • expected_dtype: A numpy dtype
  • val: An object convertible to np.array
  • message: Optional diagnostic message.

Raises:

  • ValueError: If dtype does not match.

batch_size

View source

batch_size(
    val,
    name=None
)

Returns the first (batch) dimension of val.

broadcast_to_shape_of

View source

broadcast_to_shape_of(
    val,
    target,
    name=None
)

Broadcasts val to the shape of target.

Args:

  • val: Python or Numpy array to be broadcast. Must be np.array compatible and broadcast-compatible with target.
  • target: Python or Numpy array whose shape we broadcast val to match.
  • name: Optional name for the op.

Returns:

  • broadcast_val: A np.ndarray with shape matching val + target. Provided that val's dimension sizes are all smaller or equal to target's, the returned value will be the shape of target.

cond

View source

cond(
    pred,
    true_fn,
    false_fn,
    name=None
)

Implements a conditional operation for the backend.

Args:

  • pred: A Python or Numpy bool scalar indicating the condition.
  • true_fn: A callable accepting and returning nests of np.ndarrays with the same structure as state, to be executed when pred is True.
  • false_fn: A callable accepting and returning nests of np.ndarrays with the same structure as state, to be executed when pred is False.
  • name: Optional name for the op.

Returns:

  • state: Output state, matching nest structure of input argument state.

create_variable

View source

create_variable(
    name,
    alloc,
    type_,
    max_stack_depth,
    batch_size
)

Returns an intialized Variable.

Args:

  • name: Name for the variable.
  • alloc: VariableAllocation for the variable.
  • type_: instructions.TensorType describing the sub-batch shape and dtype of the variable being created.
  • max_stack_depth: Python int, the maximum stack depth to enforce.
  • batch_size: Python int, the number of parallel threads being executed.

Returns:

  • var: A new, initialized Variable object.

equal

View source

equal(
    t1,
    t2,
    name=None
)

Implements equality comparison for Numpy backend.

fill

View source

fill(
    value,
    size,
    dtype,
    shape,
    name=None
)

Fill a fresh batched Tensor of the given shape and dtype with value.

Args:

  • value: Scalar to fill with.
  • size: Scalar int Tensor specifying the number of VM threads.
  • dtype: tf.DType of the zeros to be returned.
  • shape: Rank 1 int Tensor, the per-thread value shape.
  • name: Optional name for the op.

Returns:

  • result: Tensor of dtype values with shape [size, *shape]

full_mask

View source

full_mask(
    size,
    name=None
)

Returns an all-True mask np.ndarray with shape [size].

merge_dtypes

View source

merge_dtypes(
    dt1,
    dt2
)

Merges two dtypes, returning a compatible dtype.

Args:

  • dt1: A numpy dtype, or None.
  • dt2: A numpy dtype, or None.

Returns:

  • dtype: The more precise numpy dtype (e.g. prefers int64 over int32).

merge_shapes

View source

merge_shapes(
    s1,
    s2
)

Merges two shapes, returning a broadcasted shape.

Args:

  • s1: A list of Python int or None.
  • s2: A list of Python int or None.

Returns:

  • shape: A list of Python int or None.

Raises:

  • ValueError: If s1 and s2 are not broadcast compatible.

not_equal

View source

not_equal(
    t1,
    t2,
    name=None
)

Implements inequality comparison for Numpy backend.

prepare_for_cond

View source

prepare_for_cond(state)

Backend hook for preparing Tensors for cond.

Does nothing in the numpy backend (needed by the TensorFlow backend).

Args:

  • state: A state to be prepared for use in conditionals.

Returns:

  • state: The prepared state.

reduce_min

View source

reduce_min(
    t,
    name=None
)

Implements reduce_min for Numpy backend.

run_on_dummies

View source

run_on_dummies(
    primitive_callable,
    input_types
)

Runs the given primitive_callable with dummy input.

This is useful for examining the outputs for the purpose of type inference.

Args:

  • primitive_callable: A python callable.
  • input_types: list of instructions.Type type of each argument to the callable. Note that the contained TensorType objects must match the dimensions with which the primitive is to be invoked at runtime, even though type inference conventionally does not store the batch dimension in the TensorTypes.

Returns:

  • outputs: pattern of backend-specific objects whose types may be analyzed by the caller with type_of.

static_value

View source

static_value(t)

Gets the eager/immediate value of t.

switch_case

View source

switch_case(
    branch_selector,
    branch_callables,
    name=None
)

Implements a switch (branch_selector) { case ... } construct.

type_of

View source

type_of(
    t,
    dtype_hint=None
)

Returns the instructions.Type of t.

Args:

  • t: np.ndarray or a Python constant.
  • dtype_hint: dtype to prefer, if t is a constant.

Returns:

where

View source

where(
    condition,
    x,
    y,
    name=None
)

Implements a where selector for the Numpy backend.

Extends tf.where to support broadcasting of on_false.

Args:

  • condition: A bool np.ndarray, either a vector having length y.shape[0] or matching the full shape of y.
  • x: np.ndarray of values to take when condition is True.
  • y: np.ndarray of values to take when condition is False. May be smaller than x, as long as it is broadcast-compatible.
  • name: Optional name for the op.

Returns:

  • masked: A np.ndarray where indices corresponding to True values in condition come from the corresponding value in x, and others come from y.

while_loop

View source

while_loop(
    cond,
    body,
    loop_vars,
    name=None
)

Implements while loops for Numpy backend.

wrap_straightline_callable

View source

wrap_straightline_callable(f)

Method exists solely to be stubbed, i.e. for defun or XLA compile.