# Copyright 2022 The TensorFlow Authors. All Rights Reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# ==============================================================================

"""Utilities for slicing in to a `LinearOperator`."""

import collections

import functools

import numpy as np

from tensorflow.python.framework import dtypes

from tensorflow.python.framework import tensor_util

from tensorflow.python.ops import array_ops

from tensorflow.python.util import nest

__all__ = ['batch_slice']

def _prefer_static_where(condition, x, y):

args = [condition, x, y]

constant_args = [tensor_util.constant_value(a) for a in args]

# Do this statically.

if all(arg is not None for arg in constant_args):

condition_, x_, y_ = constant_args

return np.where(condition_, x_, y_)

return array_ops.where(condition, x, y)

def _broadcast_parameter_with_batch_shape(

param, param_ndims_to_matrix_ndims, batch_shape):

"""Broadcasts `param` with the given batch shape, recursively."""

if hasattr(param, 'batch_shape_tensor'):

# Recursively broadcast every parameter inside the operator.

override_dict = {}

for name, ndims in param._experimental_parameter_ndims_to_matrix_ndims.items(): # pylint:disable=protected-access,line-too-long

sub_param = getattr(param, name)

override_dict[name] = nest.map_structure_up_to(

sub_param, functools.partial(

_broadcast_parameter_with_batch_shape,

batch_shape=batch_shape), sub_param, ndims)

parameters = dict(param.parameters, **override_dict)

return type(param)(**parameters)

base_shape = array_ops.concat(

[batch_shape, array_ops.ones(

[param_ndims_to_matrix_ndims], dtype=dtypes.int32)], axis=0)

return array_ops.broadcast_to(

param,

array_ops.broadcast_dynamic_shape(base_shape, array_ops.shape(param)))

def _sanitize_slices(slices, intended_shape, deficient_shape):

"""Restricts slices to avoid overflowing size-1 (broadcast) dimensions.

Args:

slices: iterable of slices received by `__getitem__`.

intended_shape: int `Tensor` shape for which the slices were intended.

deficient_shape: int `Tensor` shape to which the slices will be applied.

Must have the same rank as `intended_shape`.

Returns:

sanitized_slices: Python `list` of slice objects.

"""

sanitized_slices = []

idx = 0

for slc in slices:

if slc is Ellipsis: # Switch over to negative indexing.

if idx < 0:

raise ValueError('Found multiple `...` in slices {}'.format(slices))

num_remaining_non_newaxis_slices = sum(

s is not array_ops.newaxis for s in slices[

slices.index(Ellipsis) + 1:])

idx = -num_remaining_non_newaxis_slices

elif slc is array_ops.newaxis:

pass

else:

is_broadcast = intended_shape[idx] > deficient_shape[idx]

if isinstance(slc, slice):

# Slices are denoted by start:stop:step.

start, stop, step = slc.start, slc.stop, slc.step

if start is not None:

start = _prefer_static_where(is_broadcast, 0, start)

if stop is not None:

stop = _prefer_static_where(is_broadcast, 1, stop)

if step is not None:

step = _prefer_static_where(is_broadcast, 1, step)

slc = slice(start, stop, step)

else: # int, or int Tensor, e.g. d[d.batch_shape_tensor()[0] // 2]

slc = _prefer_static_where(is_broadcast, 0, slc)

idx += 1

sanitized_slices.append(slc)

return sanitized_slices

def _slice_single_param(

param, param_ndims_to_matrix_ndims, slices, batch_shape):

"""Slices into the batch shape of a single parameter.

Args:

param: The original parameter to slice; either a `Tensor` or an object

with batch shape (LinearOperator).

param_ndims_to_matrix_ndims: `int` number of right-most dimensions used for

inferring matrix shape of the `LinearOperator`. For non-Tensor

parameters, this is the number of this param's batch dimensions used by

the matrix shape of the parent object.

slices: iterable of slices received by `__getitem__`.

batch_shape: The parameterized object's batch shape `Tensor`.

Returns:

new_param: Instance of the same type as `param`, batch-sliced according to

`slices`.

"""

# Broadcast the parammeter to have full batch rank.

param = _broadcast_parameter_with_batch_shape(

param, param_ndims_to_matrix_ndims, array_ops.ones_like(batch_shape))

if hasattr(param, 'batch_shape_tensor'):

param_batch_shape = param.batch_shape_tensor()

else:

param_batch_shape = array_ops.shape(param)

# Truncate by param_ndims_to_matrix_ndims

param_batch_rank = array_ops.size(param_batch_shape)

param_batch_shape = param_batch_shape[

:(param_batch_rank - param_ndims_to_matrix_ndims)]

# At this point the param should have full batch rank, *unless* it's an

# atomic object like `tfb.Identity()` incapable of having any batch rank.

if (tensor_util.constant_value(array_ops.size(batch_shape)) != 0 and

tensor_util.constant_value(array_ops.size(param_batch_shape)) == 0):

return param

param_slices = _sanitize_slices(

slices, intended_shape=batch_shape, deficient_shape=param_batch_shape)

# Extend `param_slices` (which represents slicing into the

# parameter's batch shape) with the parameter's event ndims. For example, if

# `params_ndims == 1`, then `[i, ..., j]` would become `[i, ..., j, :]`.

if param_ndims_to_matrix_ndims > 0:

if Ellipsis not in [

slc for slc in slices if not tensor_util.is_tensor(slc)]:

param_slices.append(Ellipsis)

param_slices += [slice(None)] * param_ndims_to_matrix_ndims

return param.__getitem__(tuple(param_slices))

def batch_slice(linop, params_overrides, slices):

"""Slices `linop` along its batch dimensions.

Args:

linop: A `LinearOperator` instance.

params_overrides: A `dict` of parameter overrides.

slices: A `slice` or `int` or `int` `Tensor` or `tf.newaxis` or `tuple`

thereof. (e.g. the argument of a `__getitem__` method).

Returns:

new_linop: A batch-sliced `LinearOperator`.

"""

if not isinstance(slices, collections.abc.Sequence):

slices = (slices,)

if len(slices) == 1 and slices[0] is Ellipsis:

override_dict = {}

else:

batch_shape = linop.batch_shape_tensor()

override_dict = {}

for param_name, param_ndims_to_matrix_ndims in linop._experimental_parameter_ndims_to_matrix_ndims.items(): # pylint:disable=protected-access,line-too-long

param = getattr(linop, param_name)

# These represent optional `Tensor` parameters.

if param is not None:

override_dict[param_name] = nest.map_structure_up_to(

param, functools.partial(

_slice_single_param, slices=slices, batch_shape=batch_shape),

param, param_ndims_to_matrix_ndims)

override_dict.update(params_overrides)

parameters = dict(linop.parameters, **override_dict)

return type(linop)(**parameters)