03-indexing-and-slicing

Select, slice, and mask — extract exactly the data you need. This example uses a grade book to demonstrate every way Nx lets you reach into an array.

dune exec nx/examples/03-indexing-and-slicing/main.exe

What You'll Learn

• Reading single elements with item
• Selecting rows and columns with I and A
• Range slicing with R and strided slicing with Rs
• Infix indexing syntax: .%{} and .${}
• Boolean masks with compress and where
• Picking rows by index with take

Key Functions

Output Walkthrough

The example starts with a 5×4 grade book (5 students, 4 subjects):

Grade book (students × subjects):
[[88, 72, 95, 83],
 [45, 90, 67, 78],
 [92, 85, 91, 70],
 [76, 63, 80, 95],
 [60, 78, 55, 82]]

Single element

item [ 0; 1 ] grades              (* → 72.0 *)

Row and column selection

The infix .${[...]} operator makes slicing readable. I n picks one index, A keeps the full axis:

grades.${[ I 2; A ]}              (* student 2, all subjects → [92, 85, 91, 70] *)
grades.${[ A; I 0 ]}              (* all students, Math      → [88, 45, 92, 76, 60] *)

Range and strided slicing

R (start, stop) is a half-open range. Rs (start, stop, step) adds a stride:

grades.${[ R (1, 4); R (0, 2) ]}  (* students 1-3, Math & Science *)
grades.${[ Rs (0, 5, 2); Rs (0, 4, 2) ]}  (* every other student & subject *)

Boolean masks

Build a boolean mask, then use compress to filter rows:

let high_math = greater_s (grades.${[ A; I 0 ]}) 85.0 in
compress ~axis:0 ~condition:high_math grades

Math > 85 mask: [true, false, true, false, false]
Students with Math > 85:
[[88, 72, 95, 83],
 [92, 85, 91, 70]]

Conditional replacement

where replaces elements based on a condition — here, flooring all grades below 60:

where (less_s grades 60.0) (full float64 [| 5; 4 |] 60.0) grades

Index Types at a Glance

Try It

1. Extract the Art column (column 3) for all students.
2. Use Rs (4, -1, -1) to reverse the student order (negative step).
3. Find students whose average grade across all subjects exceeds 80 using mean ~axes:[1] and a boolean mask.

Next Steps

Continue to 04-reshaping-and-broadcasting to learn how to change array shapes and let broadcasting align dimensions automatically.

(** Select, slice, and mask — extract exactly the data you need.

    A grade book of 5 students across 4 subjects. We'll pull out individual
    scores, entire rows and columns, ranges, and use boolean masks to find top
    performers. *)

open Nx
open Nx.Infix

let () =
  (* Grade book: 5 students × 4 subjects (Math, Science, English, Art). *)
  let grades =
    create float64 [| 5; 4 |]
      [|
        88.0;
        72.0;
        95.0;
        83.0;
        45.0;
        90.0;
        67.0;
        78.0;
        92.0;
        85.0;
        91.0;
        70.0;
        76.0;
        63.0;
        80.0;
        95.0;
        60.0;
        78.0;
        55.0;
        82.0;
      |]
  in
  Printf.printf "Grade book (students × subjects):\n%s\n\n"
    (data_to_string grades);

  (* Single element: student 0's Science score (row 0, col 1). *)
  let score = item [ 0; 1 ] grades in
  Printf.printf "Student 0, Science: %.0f\n\n" score;

  (* Entire row: all of student 2's grades. *)
  let student_2 = grades.${[ I 2; A ]} in
  Printf.printf "Student 2 (all subjects): %s\n\n" (data_to_string student_2);

  (* Entire column: everyone's Math scores (column 0). *)
  let math = grades.${[ A; I 0 ]} in
  Printf.printf "Math scores (all students): %s\n\n" (data_to_string math);

  (* Range: students 1-3, first two subjects. *)
  let subset = grades.${[ R (1, 4); R (0, 2) ]} in
  Printf.printf "Students 1-3, Math & Science:\n%s\n\n" (data_to_string subset);

  (* Strided: every other student, every other subject. *)
  let strided = grades.${[ Rs (0, 5, 2); Rs (0, 4, 2) ]} in
  Printf.printf "Every other student & subject:\n%s\n\n"
    (data_to_string strided);

  (* Boolean mask: which students scored above 85 in Math? *)
  let math_scores = grades.${[ A; I 0 ]} in
  let high_math = greater_s math_scores 85.0 in
  Printf.printf "Math > 85 mask: %s\n" (data_to_string high_math);

  let top_students = compress ~axis:0 ~condition:high_math grades in
  Printf.printf "Students with Math > 85:\n%s\n\n" (data_to_string top_students);

  (* where: replace failing grades (<60) with 60. *)
  let passing =
    where (less_s grades 60.0) (full float64 [| 5; 4 |] 60.0) grades
  in
  Printf.printf "After floor at 60:\n%s\n\n" (data_to_string passing);

  (* take: select specific students by index. *)
  let picks = take ~axis:0 (create int32 [| 3 |] [| 0l; 2l; 4l |]) grades in
  Printf.printf "Students 0, 2, 4:\n%s\n" (data_to_string picks)