Implement sampled addmm
Hi @mattbahr Can you please check @cantonios's comments and keep us posted ? Thank you! @gbaned absolutely!
@cantonios thanks for the comments! I’m going to be busy traveling for a bit but I’ll be able to get the requested changes made sometime next week 😭😭😭😭😭

On Sat, Jan 27, 2024, 1:39 AM Matt Bahr ***@***.***> wrote:

> ***@***.**** commented on this pull request.
> ------------------------------
>
> In tensorflow/python/ops/math_ops.py
> <https://github.com/tensorflow/tensorflow/pull/62750#discussion_r1468364264>
> :
>
> > +      mat1 = ops.convert_to_tensor(mat1)
> +    if not isinstance(mat2, tensor_lib.Tensor):
> +      mat2 = ops.convert_to_tensor(mat2)
> +
> +  if values.dtype != output_type:
> +    values = cast(values, output_type)
> +  if mat1.dtype != output_type:
> +    mat1 = cast(mat1, output_type)
> +  if mat2.dtype != output_type:
> +    mat2 = cast(mat2, output_type)
> +
> +  dense_rows = mat1.shape[-2]
> +  dense_cols = mat2.shape[-1]
> +
> +  # TODO(mattbahr): use dense_shape to validate the shapes of mat1 and mat2
> +  dense_shape = constant_op.constant([dense_rows, dense_cols])
>
> Importing check_ops introduces a circular dependency. I tried using a
> regular assert, but I kept getting AssertionErrors when I run my unit
> tests.
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/tensorflow/tensorflow/pull/62750#discussion_r1468364264>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/A5CIWVH2QNOI4J7IY4FRHPTYQSOJTAVCNFSM6AAAAABBQB7C26VHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMYTQNBWHA4TINBQGY>
> .
> You are receiving this because you are subscribed to this thread.Message
> ID: ***@***.***>
>
 @cantonios When you get the chance, you mind taking a look at my solution for the dense shape validation? Appreciate it! > The doc test is failing with:
> 
> ValueError: Tensor conversion requested dtype float32 for Tensor with dtype int32: <tf.Tensor: shape=(2, 3), dtype=int32, numpy= array([[1, 2, 3], [4, 5, 6]], dtype=int32)>
> 
> Please double-check your examples all run.

Hmm, my unit tests are all passing when I run on my system, but I can replicate the error if I run the following:
```
>>> indices = tf.constant([[0, 0],[1,1]])
>>> indices = tf.constant([[0, 0],[1, 1]])
>>> values = tf.constant([0.5, 0.3], dtype=tf.float32)
>>> dense_shape = tf.constant([2, 2])
>>> mat1 = tf.constant([1, 2, 3, 4, 5, 6], shape=[2, 3], dtype=tf.int32)
>>> mat2 = tf.constant([7, 8, 9, 10, 11, 12], shape=[3, 2], dtype=tf.float32)
>>> ind, res, ds = tf.sparse.sampled_addmm(indices, values, dense_shape, mat1, mat2, output_type=tf.float32)
```
@cantonios I suppose we could add a validation at the start where if `values`, `mat1`, and `mat2` are tensors, there types match `output_type`? The supported data types I had listed in the python function doc were `bfloat16`, `float16`, `float32`, and `float64`. I only set up the unit tests to run on those types. I mean to just fix the pydoc that you added - the examples in your documentation don't run.  In one of those examples, you implicitly create an integer tensor. Ah, derp thanks lol will do @cantonios I think we should be good. Recopied over the outputs and just modified to make pylint happy about line lengths. Sorry about that, and thanks for your help! Definitely still learning the project. Ok, now I understand what you meant by the doc test. I’m away for the weekend, but I’ll get this resolved Sunday night The doc test is passing now @cantonios I made a slight change for the API compatibility. When you get the chance, you mind taking another quick look at this? Thanks! Hi @cantonios Can you please review this PR ? Thank you! @gbaned @cantonios Is there anything I can do on my end to resolve the failing `feedback/copybara` check? I'm not sure what's causing it to fail. > @gbaned @cantonios Is there anything I can do on my end to resolve the failing `feedback/copybara` check? I'm not sure what's causing it to fail.

It is the API compatibility test, since this PR adds a new API entry point.  See https://github.com/tensorflow/tensorflow/tree/master/tensorflow/tools/api/tests
Relates to #56311 