Enable batch support for windowed_mean|variance

[nicolaspi]

This PR makes functions windowed_mean and windowed_variance to accept indices with batch dimensions.

Example:

x = np.array([[1, 2, 3], [1, 2, 3], [1, 2, 3]], dtype=np.float32)
low_indices = [[0, 0, 0], [1, 0, 0], [2, 2, 0]]
high_indices = [[3, 3, 3], [1, 2, 3], [3, 2, 1]]
tfp.stats.windowed_mean(x, low_indices=low_indices, high_indices=high_indices, axis=1)

Now gives:

<tf.Tensor: shape=(3, 3), dtype=float32, numpy=
array([[2. , 2. , 2. ],
       [0. , 1.5, 2. ],
       [3. , 0. , 1. ]], dtype=float32)>

Was previously failing with:

tensorflow.python.framework.errors_impl.InvalidArgumentError: required broadcastable shapes [Op:SelectV2]

[nicolaspi]

@axch I made changes in the code you authored, could you kindly have a look at this PR?
Thanks

[axch]

@nicolaspi thanks for the contribution! I am no longer an active maintainer of TFP, so I'm not really in a position to review your PR in detail (@jburnim please suggest someone?). On a quick look, though, I see a couple potential code style issues:

• Do we need the dependency on tf.experimental.numpy?
• Do we need the special case for rank-1 indices? Could we define a more uniform behavior instead?
• I'm guessing some of the shape munging already has relevant helpers defined elsewhere in TFP, but I don't remember off-hand
• TFP generally tries to handle static and dynamic TF shapes uniformly using prefer_static and tensorshape_util.

[nicolaspi]

Thanks for your feedback!

• Do we need the dependency on tf.experimental.numpy?

We need specifically the take_along_axis function that allow to gather the slices along each batch dimensions. I replaced the 'experimental' import path with from tensorflow.python.ops import numpy_ops.

• Do we need the special case for rank-1 indices? Could we define a more uniform behavior instead?

There is two motivations for this case. First, for backward compatibility, it is equivalent to the legacy non batched usage. Second, it is the only case I can think of where the broadcast is unambiguous when rank(indices) < rank(x).

• I'm guessing some of the shape munging already has relevant helpers defined elsewhere in TFP, but I don't remember off-hand

In any case, I modified the unit tests to test against non static shapes.

• TFP generally tries to handle static and dynamic TF shapes uniformly using prefer_static and tensorshape_util.

I made usage of prefer_static whenever possible.

[nicolaspi]

@jburnim can you please suggest a reviewer?
CC @axch

[SiegeLordEx]

I'll take a look at this.

[SiegeLordEx]

We'll need to do something different about take_along_axis.

1. (preferred) Somehow rewrite the logic using tf.gather/tf.gather_nd
2. Expose tf.experimental.numpy.take_along_axis in https://github.com/tensorflow/probability/tree/main/tensorflow_probability/python/internal/backend/numpy

As is, this is problematic since we really dislike using JAX_/NUMPY_MODE in library code.

[nicolaspi]

Thanks for the review!

1. I don't feel comfortable rewriting take_along_axis as it would duplicate already existing logics, I feel like it would produce unnecessary maintenance burden.
2. What about mapping tensorflow.experimental.numpy to numpy and jax.numpy backends?

[SiegeLordEx]

What is F? Why isn't it a scalar?

[nicolaspi]

Please check my comment below.

[SiegeLordEx]

I don't think this paragraph adds anything, it's just an implementation detail.

[nicolaspi]

We specify the implicit rules we uses for broadcasting. I updated the formulation.

[SiegeLordEx]

This contradicts the next paragraph, no?

In general, consider the non-batched version of this:

x shape: [N] + E
idx shape: [M]
output shape: [M] + E

The batching would introduce a batch dimension on the left of those shapes:

x shape: Bx + [N] + E
idx shape: Bi + [M]
output shape: broadcast(Bx, Bi) + [M] + E

Thus, the only broadcasting requirements are that Bx and Bi broadcast. I don't know where F came from.

[nicolaspi]

This contradicts the next paragraph, no?

Yes, I reformulated.

The batching would introduce a batch dimension on the left of those shapes:
Thus, the only broadcasting requirements are that Bx and Bi broadcast. I don't know where F came from.

Maybe the term 'batch' is not proper. This contribution adds the possibility to have the more general case where
idx shape is Bi + [M] + F. F could be seen as 'inner batch dimensions', but here 'batch' carries a different semantic than the standard machine learning one where it is represented by outer dims.

[SiegeLordEx]

These two functions are as complex as the thing we're testing. Is there any way we can write this via np.vectorize?

[nicolaspi]

I refactored using np.vectorize, but I am not sure it is easier to read.

[nicolaspi]

Hi @SiegeLordEx, I have assessed your comments, can you have a look? Thanks