Heads up: design changes for version 0.7 -- e.g. printing a tensor will require a context

[lukstafi]

The core fat-cutting refactoring: getting rid of the notion of hosted tensors, and getting rid of streams -- while preserving the current fine-grained contexts. I will also rewrite the backend scaffolding.
Conceptually (and maybe even API-wise?) printing or saving a tensor (node) will be something that contexts do.

[lukstafi]

This refactoring will also minimize user-facing API, with a new module Context handling everything via Context.t -- no user-facing device or event types. Sketch:

  module Context : sig
    type t
    type routine

    val cuda : ?device:int -> unit -> t
    val compile : t -> Assignments.t -> t * routine
    val run : t -> routine -> unit
    val print : t -> Tnode.t -> unit
  end

  (* Usage *)
  let ctx = Context.cuda () in
  let ctx, routine = Context.compile ctx assignments in
  let result = Context.execute ctx routine input

[lukstafi]

Since I plan for Context to not have event, we might lose the data-parallel examples until we introduce context tracking: val run : t -> routine -> t. This time round I will want to introduce execution context tracking only once it becomes a concrete need, e.g. restoring the data parallel examples.

[lukstafi]

This is already in progress, we lost the stream parallelism and data-parallel functionality; but old code is not refactored yet, just a reduced part of it is exposed in the new Context interface.

Summary by Claude Opus:
Major Refactoring: Simplified Context-based API for OCANNL

Summary

Today we completed a significant refactoring of OCANNL's training and
backend APIs, transitioning from a complex first-class module
parameter system to a cleaner Context-based API. This work was done in
preparation for v0.6.1, with an eye toward the more substantial v0.7
refactoring that will introduce a Universal Pool Allocator and remove
the current memory_mode system.

Background and Motivation

The conversation began with a critical assessment of OCANNL's design
compared to minimalist approaches like tinygrad. The user clarified
that OCANNL prioritizes academic innovation and expressiveness.

The key pain point identified was the cumbersome Backend module
parameter system that required verbose type annotations and
first-class modules throughout the codebase. The goal was to create a
simpler API that hides backend complexity (streams, events, device
management) while maintaining full functionality.

Key Changes Made

1. Created New Context Module (arrayjit/lib/context.ml)

We introduced a new Context module that serves as a simplified backend
interface:

• Existential wrapper pattern to hide backend-specific types
• Simple creation functions: Context.cuda(), Context.metal(),
  Context.cpu(), Context.auto()
• Node initialization tracking for safety - prevents running
  computations with missing inputs
• Global configuration support - Context.auto() respects
  OCANNL_BACKEND environment variable

Key design decisions:

• The Context tracks which tensor nodes have been initialized
• Context.run returns an updated context (functional style)
• Routines contain both the task and the context they were compiled
  with

2. Refactored train.ml API

Completely transformed the training API to use Context.t instead of
first-class Backend modules:

Before:
let forward_once (type buffer_ptr dev runner event optimize_ctx)
(module Backend : Backend with type buffer_ptr = buffer_ptr ...)
?ctx ?bindings t = ...

After:
let forward_once ?(hosted = true) ?(skip_init = false) ?reinit_all
?(bindings = IDX.empty) ctx t = ...

Changes to key functions:

• to_routine: Now takes Context.t instead of Backend module
• init_params: Simplified parameter order, uses Context's
  initialization tracking
• run_once, forward_once, update_once: All now use Context.t
• sgd_step: New function combining gradient computation and parameter
  updates

3. Migrated All Tests

Updated approximately 30+ test files from the old Backend API to the
new Context API:

Before:
let module Backend = (val Backends.fresh_backend ()) in
let backend = (module Backend : Backend with type buffer_ptr = ...) in
ignore (Train.forward_once backend tensor);

After:
let ctx = Context.auto () in
ignore (Train.forward_once ctx tensor);

Test files updated include:

• All files in test/operations/
• All files in test/einsum/
• All files in test/training/
• Utility files in bin/

4. Removed Deprecated Code

After completing the migration, we removed the With_context module
that was initially created as a transition aid. The new Context-based
functions are now the primary API.

Technical Challenges Addressed

1. Existential Types: Used GADT wrapper pattern to hide
   backend-specific types while maintaining type safety
2. Parameter Ordering: Standardized on Context as first parameter for
   consistency
3. Backward Compatibility: Briefly maintained deprecated aliases
   before removing them
4. Module Dependencies: Handled OCaml's restrictions on
   self-referential modules
5. Initialization Safety: Added tracking to prevent running
   computations with uninitialized inputs

Benefits of the New Design

1. Simpler API: No more verbose type annotations or first-class
   modules
2. Better Defaults: Context.auto() automatically selects the best
   available backend
3. Configuration Support: Respects environment variables and config
   files
4. Safety: Tracks tensor initialization to prevent runtime errors
5. Cleaner Tests: Test code is much more readable and maintainable
6. Future-Ready: Sets foundation for v0.7's Universal Pool Allocator

Impact

This refactoring touches nearly every test and example in the codebase
but maintains full functionality. The changes make OCANNL more
approachable for new users while preserving the advanced features that
enable academic research. The new Context API provides a clean
abstraction layer that will make future backend development and
optimization easier.

The work demonstrates OCANNL's evolution toward a balance of
innovation and usability, maintaining its unique features (like einsum
notation, shape inference, and dual DSLs) while improving the
developer experience.