tf.scan: add axis argument to scan along a specific axis
@LJKS,
Could you please elaborate about the mentioned Feature. Also, please specify the Use Cases for this feature which helps us to analyse. Thank you! tf.scan applies a function (arg: fn) in a stateful manner. It splits the input tensor (arg: elem) along axis 0 into n elements, where n=tf.shape(elem)[0]. It then applies the function iteratively while also updating/making use of the state to these n elements. The resulting elements are put together (stacked) into a tensor with leading dimension of equal size again. 

The feature I'm requesting simple allows for splitting (and iterating over) the tensor along a different axis than the leading axis. Let's run through a potentially typical example, with a tensor of shape=[batch_size, time_steps, features]. Now instead of scanning (tf.scan) over the leading (batch) dimension, one might want to tf.scan over the time_steps axis. Let's assume 'my_bs_ts_feat_tensor' is such a tensor: 

I can currently scan along the timesteps axis by applying: 

```
my_bs_ts_feat_tensor = tf.transpose(my_bs_ts_feat_tensor, perm=[1,0,2])
my_bs_ts_feat_tensor = tf.scan(fn=some_scan_fn, elem=my_bs_ts_feat_tensor)
my_bs_ts_feat_tensor = tf.transpose(my_bs_ts_feat_tensor, perm=[1,0,2])
```

The code above swaps the timestep-axis to the front via transpose, then scans the function along the now leading timestep axis, then transposes the timestep axis back and gets the batch-dim again as the leading dimension via transpose again. 

The above should be equal to the result of calling tf.scan(my_bs_ts_feat_tensor, axis=1) in what I propose. Or generally instead the proposed behaviour could be achieved (and cheaply implemented) as follows: 

let n_axes be the number of axes in the elem tensor (e.g. for my_bs_ts_feat_tensor of shape [batch_size, time_steps, features]: n_axes=3)
let axis the the axis argument I would like to propose in the scan function

then the proposed better scan with axis would be:

```
def tf.better_scan(fn, elems, axis):
    #transpose target axis to front
    perm_1 = [axis] + [i for i in range(n_axes) if not i==axis]
    elems = tf.transpose(elems, perm=perm_1)
    #apply the actual scan
    elems = tf.scan(fn=fn, elems=elems)
    #permute the dimensions back to where they belong, by inserting the target axis (with axis-index zero in elem) into permutations where it belongs
    perm_2 = list(range(1,n_axes)).insert(axis,0) 
    elems = tf.transpose(elems, perm=perm_2)
    return elem
```

Use-cases: 
 - Lots of stuff working with time-dimensions
 - RL applications: E.g. computing the discounted rewards-to-go, TD-errors, Advantages, GAEs, etc.
 - Literally anything else where scan would be applied, and where not by chance the correct split is the batch dimension
 - Makes some really annoying preprocessing much easier
 - Makes some comparatively hard-to-read code much easier to read Notes: 
 - I can of course also provide a more TF appropriate version (permutations via tf.range and inferring n_axes via tf.shape(tf.shape(elems))) if there is any upside with that
 - The same change of adding a axis argument as proposed for scan here would of course also be great for tf.map_fn Hi Thanks for the feature request.
As already suggested by you as the workaround to transpose and apply scan and transpose back again.
```
my_bs_ts_feat_tensor = tf.transpose(my_bs_ts_feat_tensor, perm=[1,0,2])
my_bs_ts_feat_tensor = tf.scan(fn=some_scan_fn, elem=my_bs_ts_feat_tensor)
my_bs_ts_feat_tensor = tf.transpose(my_bs_ts_feat_tensor, perm=[1,0,2])
``` 
Wouldn't it makes more sense to use the same methods in all your code implementation, rather than adding a new feature. Thanks!  Hi, 
I'm not sure I understand your answer correctly, but let me try: 
 - The workaround obviously works perfectly
 - However it makes the code much less readable and creates high risk of mistakes for users with less experience
 - E.g. When scanning along the 4th axis (i.e. axis with index 3) of a tensor with 6 axes: 
     - Perm1: [3,0,1,2,4,5]
     - Perm2: [1,2,3,0,4,5]
 - IMO the situation above is a very likely source for creating bugs, especially when code is handed over to other developers, as this is very much not 'intuitive'
 - From a usage perspective for operations which happen along a certain axis, in TF users are used to the axis argument
 
As such, with the fix being only ~4 lines of easy&quick code, and the change being non-breaking for existing code, this seems like an attractive feature to me.  Hi @LJKS!

Thanks for the request! Would you like to contribute a PR to add this functionality alongside some tests?
<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Feature Request

### Have you reproduced the bug with TF nightly?

Yes

### Source

source

### Tensorflow Version

2.11

### Custom Code

No

### OS Platform and Distribution

all

### Mobile device

-

### Python version

3.9

### Bazel version

-

### GCC/Compiler version

-

### CUDA/cuDNN version

-

### GPU model and memory

-

### Current Behaviour?

```shell
Current behaviour:
 - tf.scan 'unrolls' the argument fn along axis 0. 
Feature request: 
 - tf scan accepts the argument 'axis' (similar to many other tf operations), to specify along which axis the function is applied. The workaround is obviously simple (tf.transpose --> tf.scan --> tf.transpose; swapping the target axis to the front, applying scan, and swapping it back to where it belongs), but makes things awkward, especially in terms of readability. 

There seem to be lots of possible use-cases, but especially for working with time-axes (Batch, Time, Features), where scan should be used over the time axis this seems useful. Lots of examples e.g. in Reinforcement Learning for computing rewards-to-go, MC-state-value samples, GAE, etc.
```


### Standalone code to reproduce the issue

```shell
# e.g. the 'rewards-to-go', here to be computed on a rewards-batch of shape [batch_size, timesteps]
def rewards_to_go(rewards, gamma=0.99):
    discounted_rewards = tf.scan(lambda a, x: x + gamma * a, reverse=True, initializer=0., axis=1)
    return discounted_rewards
```


### Relevant log output

```shell
-
```
</details>