<commit_before>open Core_kernel.Std

let batch_normalization ?(decay = 0.99) t ~testing ~dims ~feature_count =
  (* TODO: use testing in a cond to not modify the beta/gamma vars in this setting. *)
  ignore testing;
  let type_ = Node.output_type t in
  let zero = Ops.const_float ~type_ (List.init feature_count ~f:(fun _ -> 0.)) in
  let one = Ops.const_float ~type_ (List.init feature_count ~f:(fun _ -> 1.)) in
  let one_minus_decay = Ops.scalar ~type_ (1. -. decay) in
  let beta = Var.create [ feature_count ] ~type_ ~init:zero in
  let gamma = Var.create [ feature_count ] ~type_ ~init:one in
  let batch_moments = Ops.moments t ~dims:(List.init dims ~f:Fn.id) in
  let beta =
    let update_beta =
      (* EWMA update. *)
      Ops.assignSub
        beta
        Ops.(one_minus_decay * (beta - batch_moments.mean))
    in
    Ops.identity ~control_inputs:[ Node.P update_beta ] beta
  in
  let gamma =
      (* EWMA update. *)
    let update_gamma =
      Ops.assignSub
        gamma
        Ops.(one_minus_decay * (gamma - batch_moments.variance))
    in
    Ops.identity ~control_inputs:[ Node.P update_gamma ] gamma
  in
  Ops.normalize t { mean = beta; variance = gamma }<commit_msg>Use ops.cond in the batch normalization layer.<commit_after>open Core_kernel.Std

let batch_normalization ?(decay = 0.99) t ~testing ~dims ~feature_count =
  let type_ = Node.output_type t in
  let zero = Ops.const_float ~type_ (List.init feature_count ~f:(fun _ -> 0.)) in
  let one = Ops.const_float ~type_ (List.init feature_count ~f:(fun _ -> 1.)) in
  let one_minus_decay = Ops.scalar ~type_ (1. -. decay) in
  let beta = Var.create [ feature_count ] ~type_ ~init:zero in
  let gamma = Var.create [ feature_count ] ~type_ ~init:one in
  let batch_moments = Ops.moments t ~dims:(List.init dims ~f:Fn.id) in
  let beta_with_update =
    let update_beta =
      (* EWMA update. *)
      Ops.assignSub
        beta
        Ops.(one_minus_decay * (beta - batch_moments.mean))
    in
    Ops.identity ~control_inputs:[ Node.P update_beta ] beta
  in
  let gamma_with_update =
      (* EWMA update. *)
    let update_gamma =
      Ops.assignSub
        gamma
        Ops.(one_minus_decay * (gamma - batch_moments.variance))
    in
    Ops.identity ~control_inputs:[ Node.P update_gamma ] gamma
  in
  let beta = Ops.cond testing ~if_true:beta ~if_false:beta_with_update in
  let gamma = Ops.cond testing ~if_true:gamma ~if_false:gamma_with_update in
  Ops.normalize t { mean = beta; variance = gamma }
{'fraction_non_alphanumeric': 0.10741482965931863, 'fraction_numerical': 0.004809619238476954, 'mean_word_length': 3.0651465798045603, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, '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': '2803866', 'n_tokens_mistral': 909, 'n_tokens_neox': 854, 'n_words': 245}