Add dedicated `TensorGenerator` subclasses for #437/`ReadDataFallback` ops

[khatchad]

Summary

ReadDataFallback (added by #196 for #437) classifies XML-modeled tensor-producing APIs that don't have a dedicated TensorGenerator subclass yet as a generic ⊤-shape / UNKNOWN-dtype tensor source. That's enough to satisfy "is this a tensor source?" checks (e.g. Hybridize's getHasTensorParameter()) but loses precision: the ops it covers all do have well-defined output shapes and dtypes the analysis could be capturing.

This issue tracks adding per-op TensorGenerator subclasses with proper precision, so the fallback's ⊤ result becomes the exception rather than the norm. Once every op in scope has its own subclass, ReadDataFallback itself can be removed.

Ops Currently Routed Through ReadDataFallback

These are the ops surfaced by wala/ML#437's investigation that hit the fallback (each hits the analysis but with ⊤ shape / UNKNOWN dtype):

Op	Output shape (semantic)	Output dtype
tf.rank(t)	() (scalar)	int32
tf.stack(values, axis)	t.shape[:axis] + (len(values),) + t.shape[axis:]	inherits from values[0]
tf.linspace(start, stop, num)	(num,)	float32 (or start.dtype if specified)
tf.range(start, limit, delta)	derived from start/limit/delta	inherits from args (already has Range in dispatch — partial precision)
tf.boolean_mask(tensor, mask)	depends on mask	inherits from tensor
tf.einsum(equation, *inputs)	derived from equation	inherits from inputs
tf.math.top_k(input, k)	tuple of (values, indices), both shape input.shape[:-1] + (k,)	values: input.dtype; indices: int32
tf.meshgrid(*xi)	tuple of len(xi) tensors	inherits from xi[0]
tf.image.extract_patches(images, sizes, strides, rates, padding)	derived	inherits from images
tf.math.exp(x) / tf.math.exp2(x)	x.shape	x.dtype
tf.math.argmax(input, axis)	input.shape minus axis	int64 (default) or int32
tf.math.argmin(input, axis)	same as argmax	same
... (full list per wala/ML#380's 22-op enumeration)

Approach

For each op:

1. Add a subclass <Op> extends TensorGenerator (or extends TensorGeneratorWithSingleInput / similar shared base if patterns emerge).
2. Override getDefaultShapes to compute the precise output shape from the input parameter(s).
3. Override getDefaultDTypes to return the right dtype (often inherited from input, sometimes fixed like int32/int64 for argmax/rank).
4. Add a dispatch entry: else if (isType(calledFunction, OP.getDeclaringClass())) return new <Op>(source);.
5. Add an Ariadne-side test: tf2_test_<op>.py + TestTensorflow2Model.test<Op> asserting the precise shape/dtype (per the test-fixture three-check protocol).
6. Once the subclass is in place, the ReadDataFallback no longer fires for that op.

Acceptance Criteria

• All ops in the table above have dedicated TensorGenerator subclasses in com.ibm.wala.cast.python.ml.client/.
• Each op has Ariadne-side tests asserting the precise shape/dtype.
• ReadDataFallback no longer fires for any test in Ariadne's com.ibm.wala.cast.python.ml.test suite (verifiable by running the suite with LOGGER.fine enabled and checking for read_data-pattern fallback log lines).
• ReadDataFallback, getPropertyReadMemberNames, and getTensorflowReadDataPropertyNames removed from TensorGeneratorFactory.

Cross-Refs

• Tensor classification regresses on tf.rank/tf.stack/tf.linspace/tf.range/etc. — TensorTypeAnalysis lacks entries despite unchanged XML #437 — the regression that motivated ReadDataFallback.
• Fix reshape() propagation #196 (ponder-lab/ML#196) — adds ReadDataFallback; this issue tracks its replacement.
• Systematically inline remaining read_data/read_dataset marker allocations in tensorflow.xml #380 — separate but related: inlining read_data XML markers (the analyzer-side work here is the next step after).
• Design issues in TensorGeneratorFactory.getFunction() and dispatchByPropertyName exposed by #437 #448 — getFunction() design refactor; orthogonal but would let the dispatch/fallback boundary be cleaner.

[khatchad]

TDD Scaffolding Plan (Deferred — Resume After ponder-lab/Hybridize-Functions-Refactoring#427 Merges)

To make this issue's per-op work TDD-driven and visible, the proposed scaffolding step is a draft PR on ponder-lab/ML (fork-only, not upstream) containing only failing tests — one per op in scope. Each generator that lands flips its corresponding test from red to green. The PR shape is then "remove ignore + add <Op> extends TensorGenerator + add dispatch entry + close test fixture's runtime asserts." The remaining-red count = remaining work.

Ops in Scope

Per the table in this issue's body — these are the ops currently routed through ReadDataFallback and need dedicated subclasses:

• tf.rank
• tf.stack
• tf.linspace
• tf.boolean_mask
• tf.einsum
• tf.meshgrid
• tf.math.exp / tf.math.exp2
• tf.math.top_k
• tf.math.argmax / tf.math.argmin
• tf.image.extract_patches
• tf.range (already has a Range generator with partial precision; tighten dtype handling for the iterating-element case described in Tensor classification regresses on tf.rank/tf.stack/tf.linspace/tf.range/etc. — TensorTypeAnalysis lacks entries despite unchanged XML #437)

Scaffold Shape (Per Op)

For op tf.<op> (e.g., tf.rank):

1. Python fixture at com.ibm.wala.cast.python.test/data/tf2_test_<op>.py:

   import tensorflow as tf
   def f(a):
       pass
   t = tf.constant([[[1, 1, 1], [2, 2, 2]], [[3, 3, 3], [4, 4, 4]]])
   r = tf.<op>(t)  # adapted per-op
   assert isinstance(r, tf.Tensor)
   assert r.dtype == tf.<expected_dtype>
   assert r.shape == <expected_shape>
   f(r)

2. JUnit test in TestTensorflow2Model:

   @Test
   public void test<Op>() throws ... {
     test(\"tf2_test_<op>.py\", \"f\", 1, 1, Map.of(2, Set.of(<concrete TensorType>)));
   }

3. Initially fails because ReadDataFallback produces ⊤-shape/UNKNOWN-dtype, not the asserted concrete TensorType. The implementer adds a <Op> generator subclass + dispatch entry; the test goes green.

Why This Was Deferred

Hybridize tests assert tensor identification (`getHasTensorParameter()` / `getHasPrimitiveParameter()`-style binary checks), not concrete shapes/dtypes. The scaffolding here adds precision — which is orthogonal to driving down ponder-lab/Hybridize-Functions-Refactoring#427's failure count. That failure count is dominated by other mechanisms (recursion, retracing, side-effects, range-iteration). After ponder-lab/Hybridize-Functions-Refactoring#427 merges, this scaffolding becomes the next work surface.

Execution Order Once Resumed

1. Open the draft PR with the full scaffold (one fixture + one ignored/failing test per op).
2. Pick the simplest op (probably tf.rank — scalar int32 output, no input dependencies — already drafted earlier).
3. Land per-op generators one PR each, each flipping a single test from red to green.
4. When all are green, this issue can close; ReadDataFallback and its trampoline-class scan can be retired.

🤖 Generated with Claude Code

[khatchad]

Post-#197 triage: 45 unhandled ops in the superset

Updated count after the tensorflow.xml superset sync (#197). Verified via diff of tensorflow.xml (50 new <class> declarations, minus 5 non-op aliases like Custom_gradient / TPUStrategy / OneDeviceStrategy / Input / shuffle which already have other handling). All 45 currently fall through to ReadDataFallback — same ⊤-shape / UNKNOWN-dtype outcome — and need dedicated generators per this issue's scope.

Grouped by implementation complexity so the easy wins go first:

Tier 1 — pure pass-through (shape + dtype from arg 0)

Trivially mechanical; mirrors the Sigmoid / Softmax precedent. Each is ~30 LOC of generator + ~10 LOC of factory dispatch + a 5-line Python fixture + a one-line test method.

Op	Notes
tf.identity(input)	Literally pass-through. Reference impl.
tf.stop_gradient(input)	Pass-through; only differs from identity in autograd semantics, irrelevant to static analysis.
tf.bias_add(value, bias)	Shape + dtype from value. bias is broadcast-added but doesn't change the receiver's shape.
tf.clip_by_global_norm(t_list, ...)	Returns (t_list, global_norm). Each element pass-through; global_norm is ()-shape float32.

Tier 2 — element-wise unary math (shape from input, dtype from input)

Same pattern as Sigmoid's getDefaultShapes/getDefaultDTypes (read input arg, return its shape + dtype). Five-line copy-paste-rename per op.

Op	Constraints
tf.math.abs(x)	Any numeric dtype
tf.math.acos(x)	Float dtype required at runtime; static analysis can preserve input's dtype
tf.math.exp(x)	Float dtype
tf.math.tanh(x)	Float dtype
tf.math.rsqrt(x)	Float dtype
tf.nn.log_softmax(logits)	Same shape, float dtype preserved
tf.math.l2_normalize(x, axis)	Shape preserved (axis collapses, but normalize keeps dimensions)

Tier 3 — reductions (dims collapse along axis, dtype-preserving)

Same pattern as ReduceSum / ReduceMean. The shape rewrite (drop axis dim or keep with keepdims=True) is already factored in those existing generators — just subclass and reuse.

Op
tf.math.reduce_all(input_tensor, axis, keepdims)
tf.math.reduce_max(input_tensor, axis, keepdims)
tf.math.reduce_prod(input_tensor, axis, keepdims)
tf.math.reduce_logsumexp(input_tensor, axis, keepdims)

Tier 4 — fixed output shape and dtype

Hardcoded () scalar + int32 — same shape as DatasetRangeGenerator's intrinsic-API pattern.

Op	Output
tf.rank(input)	() int32
tf.size(input)	() int32

Tier 5 — multi-input combinators

Output shape derives from the list of input tensors. Reasonably mechanical but each has its own shape-derivation rule.

Op	Shape rule
tf.add_n(inputs)	inputs[0].shape (all required to match); dtype from inputs[0]
tf.concat(values, axis)	values[0].shape with axis dim summed across values
tf.stack(values, axis)	values[0].shape with new len(values) dim inserted at axis
tf.meshgrid(*xi)	Tuple of len(xi) tensors, each xi[0].shape × ... × xi[-1].shape. Tuple semantics — needs TupleElementProvider.

Tier 6 — shape transformations from args (heterogeneous, harder)

Each has a non-trivial shape rule. Implement individually with care.

tf.linspace(start, stop, num), tf.math.top_k(input, k) (tuple result), tf.tensordot(a, b, axes), tf.einsum(equation, *inputs) (equation parsing!), tf.gather_nd(params, indices), tf.nn.embedding_lookup(params, ids), tf.dynamic_partition(data, partitions, num_partitions) (returns list), tf.dynamic_stitch(indices, data), tf.image.extract_patches(images, sizes, strides, rates, padding), tf.tensor_scatter_nd_update(tensor, indices, updates), tf.sequence_mask(lengths, maxlen, dtype), tf.where(condition, x, y), tf.split(value, num_or_size_splits, axis) (returns list), tf.unstack(value, num, axis) (returns list), tf.sort(values, axis), tf.linalg.eigh(tensor) (tuple result), tf.linalg.trace(x).

Tier 7 — control flow (output depends on body, hardest)

Need to walk the function body (true_fn/false_fn/body) and resolve the output shape/dtype from there. Significant infra likely needed.

tf.cond(pred, true_fn, false_fn), tf.while_loop(cond, body, loop_vars), tf.map_fn(fn, elems), tf.experimental_run_v2(...), tf.py_function(func, inp, Tout) (Tout is the explicit output dtype — actually a Tier 5/6 candidate since it's structurally simpler than the others here), tf.custom_gradient(f).

Suggested order

1. Land Tier 1 + Tier 2 + Tier 4 first (~13 ops; truly mechanical; one PR per tier or one big PR with all 13).
2. Tier 3 next — reuses ReduceSum/ReduceMean shape-rewriting infra; another PR.
3. Tier 5 with care for tuple semantics on meshgrid.
4. Tier 6 individually as needed; prioritize whichever ops show up in concrete fixtures from wala/ML#437's test corpus.
5. Tier 7 last; needs design discussion.

The acceptance criterion of "removes ReadDataFallback" can land at the end of Tier 6 once every fallback-routed op has a dedicated generator.