System information

• Version: master
• Are you willing to contribute it: Yes

Describe the feature and the current behavior/state.

There's a good amount of ops that take indices, shapes, permutations, scalars, or some kind of often constant value as inputs. Currently, these have to be converted to operand, which results in lots of wrapping and prevents the use of defaults (things like reducing all dimensions if none are specified, for reduce ops). Examples include all the reduction ops (mean, all, sum), reshape, and transpose (permute),

Will this change the current api? How?

I'd like to create a WrapperHelpers class, like Helpers, to create Ops methods that take Java arrays, or varargs where possible, for the most common of these Ops.

For example:

// Current API
tf.math.mean(x, tf.array(0, 1, 2));
// With wrappers
tf.math.mean(x, new long{0, 1, 2});
// OR, if it's a rank 3 array
tf.math.mean(x, null);

Who will benefit with this feature?

Anyone using these ops (and they are common ops). This isn't a huge improvement in size (the example is actually longer), the biggest improvement is from being able to use defaults for empty arrays or nulls. We can do build-time Java-side error checking for things like repeated indices. For rank-dependent inputs like for transpose this is especially useful.

It also helps out a Kotlin API a bit, since we aren't using the Options vararg and so can use it for this. Something like tf.math.mean(x, 2, 3, keepDims = true) or x.mean(2, 3, keepDims = true) eventually.

Any Other info.

I experimented with codegen for this, as these inputs are marked with Tidx or similar in ops.proto, however many different types of indices or shapes are used (i.e. 2D for BatchToSpace or 1D for reduce) with no distinction, so I don't think it's possible to generate these wrappers.

If this goes well, I'd also like to look at creating similar methods to wrap scalars for math ops, i.e. x.add(2), since I suspect they will be commonly used.