Skip to content

Recognize tf.SparseTensorSpec on the input side so sparse-spec parameters are typed #589

Description

@khatchad

Summary

Ariadne can read a tf.TensorSpec / tf.RaggedTensorSpec that appears in analyzed source (e.g., a developer-written @tf.function(input_signature=[tf.RaggedTensorSpec(...)]), or a type_spec= argument) and recover the parameter's shape and dtype from the spec's fields. There is no tf.SparseTensorSpec analog, so a sparse spec in source is not recognized and the parameter it describes gets no tensor type.

Where the asymmetry lives

TensorFlowTypes models the readable specs:

  • TENSOR_SPEC / RAGGED_TENSOR_SPEC TypeReferences (TensorFlowTypes.java:278, :290)
  • their shape/dtype field refs — SPEC_SHAPE/SPEC_DTYPE, RAGGED_SPEC_SHAPE/RAGGED_SPEC_DTYPE (:296-306)
  • entries in the type-to-synthetic-signature map (:1834-1835)

TensorGenerator then reads them: when a value is a TensorSpec/RaggedTensorSpec, it extracts shape and dtype and infers the type (TensorGenerator.java:514, :1513).

There is no SPARSE_TENSOR_SPEC TypeReference, no SPARSE_SPEC_* field refs, and no map entry, so tf.SparseTensorSpec(shape=..., dtype=...) in source flows through untyped.

Relationship to #588

This is the input-side mirror of #588. #588 adds output-side sparse marking — tagging the result of sparse-producing ops (tf.SparseTensor, tf.sparse.add, …) so a consumer can tell sparse from dense. This issue is the complementary input-side recognition: reading a SparseTensorSpec that a developer already wrote. They are independent and both needed for full sparse parity.

Sketch

Mirror the RAGGED_TENSOR_SPEC path: add SPARSE_TENSOR_SPEC TypeReference + SPARSE_SPEC_SHAPE/SPARSE_SPEC_DTYPE field refs + the map entry, and extend the TensorGenerator spec-reading arms to handle it (marking the recovered type sparse once #588's marker exists).

Metadata

Metadata

Assignees

No one assigned

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions