Skip to content

Conversation

@copybara-service
Copy link
Contributor

Split large fully-connected nodes along k.

This is just a subgraph rewrite that replaces nodes that look like

fully_connected(input_id, weights_id, bias_id)

where the data along the kth dimension exceeds 16kB, with

fully_connected(static_slice(K0, input_id), static_slice(K0, weights_id), bias_id)
  +
fully_connected(static_slice(K1, input_id), static_slice(K1, weights_id), 0)
  +
...
  +
fully_connected(static_slice(Kn, input_id), static_slice(Kn, weights_id), 0)

where static_slice(ival, value_id) slices value_id in the k dimension.

@copybara-service copybara-service bot force-pushed the test_806264392 branch 2 times, most recently from 3dbc2f2 to b7a9727 Compare November 6, 2025 16:13
This is just a `subgraph` rewrite that replaces nodes that look like
```
fully_connected(input_id, weights_id, bias_id)
```
where the data along the `k`th dimension exceeds 16kB, with
```
fully_connected(static_slice(K0, input_id), static_slice(K0, weights_id), bias_id)
  +
fully_connected(static_slice(K1, input_id), static_slice(K1, weights_id), 0)
  +
...
  +
fully_connected(static_slice(Kn, input_id), static_slice(Kn, weights_id), 0)
```
where `static_slice(ival, value_id)` slices `value_id` in the `k` dimension.

PiperOrigin-RevId: 806264392
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant