src/test/scripts/applications/nn/component/shampoo_test2.dml

#-------------------------------------------------------------
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#
#-------------------------------------------------------------

source("scripts/nn/optim/shampoo.dml") as shampoo

X_main = matrix("0.12 -0.45  0.33  0.08 -0.19 -0.27  0.41 -0.05  0.22  0.14  0.09 -0.31  0.26 -0.48  0.37 0.44  0.06 -0.29  0.15 -0.11 -0.38  0.24  0.17 -0.07  0.52",
  rows=5, cols=5
)

dX_main = matrix("0.015 -0.022  0.008  0.031 -0.012 -0.009  0.027 -0.014  0.005  0.019 0.021 -0.006  0.011 -0.025  0.004 -0.018  0.013 -0.029  0.007 -0.016 0.010 -0.017  0.024 -0.003  0.028",
  rows=5, cols=5
)

epsilon = 1e-4
lr = 0.005
diagThreshold = 10
rootEvery = 10
preconEvery = 10

for (diagThreshold in seq(1, 10, 9)){
  X = X_main
  dX = dX_main
  #[preconL, preconR, useDiag] = shampoo::init(X, epsilon, diagThreshold)
  #[X, preconL, preconR] = shampoo::update(X, dX, lr, preconL, preconR, useDiag)

  #[preconL, preconR, momentum, useDiag] = shampoo::init_momentum(X, epsilon, diagThreshold)
  #[X, preconL, preconR, momentum] = shampoo::update_momentum(X, dX, lr, preconL, preconR, momentum, useDiag)

  [preconL, preconR, stepCounter, bufferL, bufferR, momentum, preconLInvPowerRoot, preconRInvPowerRoot, useDiag] = shampoo::init_heuristic(X, epsilon, diagThreshold)
  [X, preconL, preconR, momentum, stepCounter, bufferL, bufferR, preconLInvPowerRoot, preconRInvPowerRoot] = shampoo::update_heuristic(X, dX, lr, preconL, preconR, momentum, stepCounter, rootEvery, preconEvery, bufferL, bufferR, preconLInvPowerRoot, preconRInvPowerRoot, useDiag)
  
  print("diagThreshold: " + diagThreshold)
  print(X)
}