<commit_before>open Lacaml_D

let scale_by_col_sum a =
  for j = 1 to Mat.dim2 a do
    let c = Mat.col a j in
    ignore (Vec.exp ~y:c c);
    let s = 1. /. Vec.sum c in
    scal s c
  done

(* Classes in y = {1,2 ... k}
   General softmax method.
*)
let general_eval_and_grad ~lambda k x y =
  let open Bigarray in
  let m = Mat.dim1 x in         (* number of data points *)
  let n = Mat.dim2 x in         (* number of features *)
  let n_y = Array.length y in
  if n_y <> m then
    invalid_arg (Printf.sprintf "different dimensions %d vs %d" n n_y)
  else
    let nmf = -1. /. float m in
    let ld2 = lambda /. 2. in
    let ind = Mat.init_cols k m (fun r c -> if r = y.(c-1) then 1. else 0.) in
    begin fun w_c (g_c : Lacaml_float64.vec) ->
      (* unroll w and g *)
      let w = reshape_2 (genarray_of_array1 w_c) k n in
      let g = reshape_2 (genarray_of_array1 g_c) k n in
      let p = gemm w ~transb:`T x in
      scale_by_col_sum p;
      let s =
        (* one could do this entirely in Fortran by
           1. convert p to a vector -> p_v
           2. Vec.log of p_v
           3. convert ind to a vector -> ind_v
           4. Vec.add p_v ind_v *)
        Array.fold_left (fun (s, j) i ->
            let v = (Mat.col p j).{i} in
            s +. log v, j + 1)
          (0.,1) y
        |> fst
        |> fun s -> nmf *. s +. ld2 *. Vec.sqr_nrm2 w_c
      in
      let d =
        Array.init m (fun j ->
            let j = j + 1 in
            Vec.sub (Mat.col ind j) (Mat.col p j))<commit_msg>Subtract the max, per column, to avoid large values<commit_after>open Lacaml_D

let remove_column_max a =
  let one = Vec.make (Mat.dim1 a) 1. in
  for j = 1 to Mat.dim2 a do
    let c = Mat.col a j in
    let mx = Vec.max c in
    axpy ~alpha:(-1. *. mx) one c
  done

let scale_by_col_sum a =
  for j = 1 to Mat.dim2 a do
    let c = Mat.col a j in
    ignore (Vec.exp ~y:c c);
    let s = 1. /. Vec.sum c in
    scal s c
  done

(* Classes in y = {1,2 ... k}
   General softmax method.
*)
let general_eval_and_grad ~lambda k x y =
  let open Bigarray in
  let m = Mat.dim1 x in         (* number of data points *)
  let n = Mat.dim2 x in         (* number of features *)
  let n_y = Array.length y in
  if n_y <> m then
    invalid_arg (Printf.sprintf "different dimensions %d vs %d" n n_y)
  else
    let nmf = -1. /. float m in
    let ld2 = lambda /. 2. in
    let ind = Mat.init_cols k m (fun r c -> if r = y.(c-1) then 1. else 0.) in
    begin fun w_c (g_c : Lacaml_float64.vec) ->
      (* unroll w and g *)
      let w = reshape_2 (genarray_of_array1 w_c) k n in
      let g = reshape_2 (genarray_of_array1 g_c) k n in
      let p = gemm w ~transb:`T x in
      remove_column_max p;
      scale_by_col_sum p;
      let s =
        (* one could do this entirely in Fortran by
           1. convert p to a vector -> p_v
           2. Vec.log of p_v
           3. convert ind to a vector -> ind_v
           4. Vec.add p_v ind_v *)
        Array.fold_left (fun (s, j) i ->
            let v = (Mat.col p j).{i} in
            s +. log v, j + 1)
          (0.,1) y
        |> fst
        |> fun s -> nmf *. s +. ld2 *. Vec.sqr_nrm2 w_c
      in
      let d =
        Array.init m (fun j ->
            let j = j + 1 in
            Vec.sub (Mat.col ind j) (Mat.col p j))
{'fraction_non_alphanumeric': 0.10743550834597876, 'fraction_numerical': 0.019119878603945373, 'mean_word_length': 1.601420678768745, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7195233', 'n_tokens_mistral': 1284, 'n_tokens_neox': 1248, 'n_words': 491}