input: Write a python function to optimize a given cost function using gradient descent or Newton's method.
  assert np.allclose(optimize(lambda x: x**2, 1), 0, atol=1e-2)
  assert np.allclose(optimize(lambda x: x**2, 2), [0, 0], atol=1e-2)
  assert np.allclose(optimize(lambda x: np.sum(x**2), 3), [0, 0, 0], atol=1e-2)
output: import numpy as np
  
  def optimize(cost_function, dimensions, method='gradient_descent', learning_rate=0.01, tol=1e-6, max_iter=1000):
      def gradient_descent(f, x0, lr, tol, max_iter):
          x = x0
          for _ in range(max_iter):
              grad = np.gradient(f(x))
              new_x = x - lr * grad
              if np.linalg.norm(new_x - x) < tol:
                  break
              x = new_x
          return x
  
      def newton_method(f, x0, tol, max_iter):
          x = x0
          for _ in range(max_iter):
              grad = np.gradient(f(x))
              hessian = np.gradient(np.gradient(f(x)))
              new_x = x - np.linalg.inv(hessian).dot(grad)
              if np.linalg.norm(new_x - x) < tol:
                  break
              x = new_x
          return x
  
      initial_guess = np.zeros(dimensions)
  
      if method == 'gradient_descent':
          return gradient_descent(cost_function, initial_guess, learning_rate, tol, max_iter)
      elif method == 'newton':
          return newton_method(cost_function, initial_guess, tol, max_iter)
      else:
          raise ValueError("Unknown optimization method")