Skip to content

Basic Nx.block implementation#1709

Draft
Chapaman wants to merge 3 commits intoelixir-nx:mainfrom
Chapaman:block_implementation
Draft

Basic Nx.block implementation#1709
Chapaman wants to merge 3 commits intoelixir-nx:mainfrom
Chapaman:block_implementation

Conversation

@Chapaman
Copy link
Contributor

Introduces Nx.block/3 as the primary Nx-side API for extensible, named computation blocks with a default implementation and backend bridges. Migrates logical_not/1 and phase/1 from Nx.Shared.optional/4 to block/3.

Changes

  • Nx.block/3 (nx/lib/nx.ex): Resolves a common backend from flattened tensor leaves in the container (supports lists/other containers); dispatches in order:
    1 - backend.block/3 if exported
    2 - backend.optional/3 with a struct-derived callback name
    3 - default fun.(container, struct)
  • struct_to_optional_callback_name/1 (nx/lib/nx.ex): Maps block structs to existing optional op atoms (e.g. LogicalNot:logical_not, Phase:phase).
  • Block structs (nx/lib/nx/block.ex): Nx.Block.LogicalNot, Nx.Block.Phase (empty structs for identity/static config for now).
  • Nx.logical_not/1: Uses block(%Nx.Block.LogicalNot{}, tensor, ...) instead of Nx.Shared.optional/4.
  • Nx.phase/1: Uses block(%Nx.Block.Phase{}, tensor, ...) instead of Nx.Shared.optional/4.

Discussion

Nx.logical_not/1 and Nx.phase/1 have been altered to showcase Nx.block/3.

  • Should we soft-deprecate Nx.Backend.optional/3 (and related patterns) in docs first, keep behavior, and migrate call sites gradually?
  • Should we have an horizontal change across the ecosystem (Nx + EXLA + Torchx, and any other backends that you see fit) in one breaking release, replacing optional-based dispatch with block (and/or a formal @callback block/3) everywhere?

@Chapaman Chapaman closed this Mar 20, 2026
@Chapaman Chapaman reopened this Mar 20, 2026
@josevalim
Copy link
Contributor

Should we soft-deprecate Nx.Backend.optional/3 (and related patterns) in docs first, keep behavior, and migrate call sites gradually?

It is already an optional callback, so I would simply stop calling it and remove it from the behaviour description.

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.

2 participants